Skip to content

Commit

Permalink
add comparison table
Browse files Browse the repository at this point in the history
  • Loading branch information
allenanie committed Oct 31, 2024
1 parent 148e7f6 commit 41be285
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 31 deletions.
114 changes: 83 additions & 31 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,38 +9,43 @@
![GitHub license](https://img.shields.io/badge/License-MIT-blue.svg)
[![arXiv](https://img.shields.io/badge/arXiv-1234.56789-b31b1b.svg)](https://arxiv.org/abs/2406.16218)



Trace is a new AutoDiff-like tool for training AI systems end-to-end with general feedback (like numerical rewards or losses, natural language text, compiler errors, etc.). Trace generalizes the back-propagation algorithm by capturing and propagating an AI system's execution trace. Trace is implemented as a PyTorch-like Python library. Users write Python code directly and can use Trace primitives to optimize certain parts, just like training neural networks!
Trace is a new AutoDiff-like tool for training AI systems end-to-end with general feedback (like numerical rewards or
losses, natural language text, compiler errors, etc.). Trace generalizes the back-propagation algorithm by capturing and
propagating an AI system's execution trace. Trace is implemented as a PyTorch-like Python library. Users write Python
code directly and can use Trace primitives to optimize certain parts, just like training neural networks!

[Paper](https://arxiv.org/abs/2406.16218) | [Project website](https://microsoft.github.io/Trace/) | [Documentation](https://microsoft.github.io/Trace/intro.html) | [Blogpost](https://www.microsoft.com/en-us/research/blog/tracing-the-path-to-self-adapting-ai-agents/)


## Setup

Simply run
Simply run

pip install trace-opt

Or for development, clone the repo and run the following.
Or for development, clone the repo and run the following.

pip install -e .

The library requires Python >= 3.9. The installation script will git clone [AutoGen](https://github.com/microsoft/autogen). You may require [Git Large File Storage](https://git-lfs.com/) if git is unable to clone the repository otherwise.
The library requires Python >= 3.9. The installation script will git
clone [AutoGen](https://github.com/microsoft/autogen). You may require [Git Large File Storage](https://git-lfs.com/) if
git is unable to clone the repository otherwise.

## Updates

- **2024.10.21** New [paper](https://arxiv.org/abs/2410.15625) by Nvidia, Stanford, Visa, & Intel applies Trace to optimize for mapper code of parallel programming. Trace (OptoPrime) learns code achieving 1.3X speed up under 10 minutes, compared with the code optimized by domain expert.
- **2024.9.25** [Trace Paper](https://arxiv.org/abs/2406.16218) is accepted to NeurIPS 2024.
- **2024.9.14** TextGrad is available as an optimizer in Trace.

- **2024.10.21** New [paper](https://arxiv.org/abs/2410.15625) by Nvidia, Stanford, Visa, & Intel applies Trace to
optimize for mapper code of parallel programming. Trace (OptoPrime) learns code achieving 1.3X speed up under 10
minutes, compared with the code optimized by domain expert.
- **2024.9.25** [Trace Paper](https://arxiv.org/abs/2406.16218) is accepted to NeurIPS 2024.
- **2024.9.14** TextGrad is available as an optimizer in Trace.

## QuickStart

Trace has two primitives: `node` and `bundle`. `node` is a primitive to define a node in the computation graph. `bundle` is a primitive to define a function that can be optimized.
Trace has two primitives: `node` and `bundle`. `node` is a primitive to define a node in the computation graph. `bundle`
is a primitive to define a function that can be optimized.

```python
from opto.trace import node

x = node(1, trainable=True)
y = node(3)
z = x / y
Expand All @@ -53,14 +58,17 @@ node_of_list.append(3)
```

Once a node is declared, all the following operations on the node object will be automatically traced.
We built many magic functions to make a node object act like a normal Python object. By marking `trainable=True`, we tell our optimizer that this node's content
We built many magic functions to make a node object act like a normal Python object. By marking `trainable=True`, we
tell our optimizer that this node's content
can be changed by the optimizer.

For functions, Trace uses decorators like @bundle to wrap over Python functions. A bundled function behaves like any other Python functions.
For functions, Trace uses decorators like @bundle to wrap over Python functions. A bundled function behaves like any
other Python functions.

```python
from opto.trace import node, bundle


@bundle(trainable=True)
def strange_sort_list(lst):
'''
Expand All @@ -71,17 +79,20 @@ def strange_sort_list(lst):
lst = sorted(lst)
return lst


test_input = [1, 2, 3, 4]
test_output = strange_sort_list(test_input)
print(test_output)
```

Now, after declaring what is trainable and what isn't, and use `node` and `bundle` to define the computation graph, we can use the optimizer to optimize the computation graph.
Now, after declaring what is trainable and what isn't, and use `node` and `bundle` to define the computation graph, we
can use the optimizer to optimize the computation graph.

```python
import autogen
from opto.optimizers import OptoPrime


# we first declare a feedback function
# think of this as the reward function (or loss function)
def get_feedback(predict, target):
Expand All @@ -90,6 +101,7 @@ def get_feedback(predict, target):
else:
return "test case failed!"


test_ground_truth = [1, 4, 2, 3]
test_input = [1, 2, 3, 4]

Expand Down Expand Up @@ -117,21 +129,50 @@ Then, we can use the familiar PyTorch-like syntax to conduct the optimization.
## Supported Optimizers

Currently, we support three optimizers:

- OPRO: [Large Language Models as Optimizers](https://arxiv.org/abs/2309.03409)
- TextGrad: [TextGrad: Automatic "Differentiation" via Text](https://arxiv.org/abs/2406.07496)
- OptoPrime: Our proposed algorithm -- using the entire computational graph to perform parameter update. It is 2-3x faster than TextGrad.
- OptoPrime: Our proposed algorithm -- using the entire computational graph to perform parameter update. It is 2-3x
faster than TextGrad.

Using our framework, you can seamlessly switch between different optimizers:

```python
optimizer1 = OptoPrime(strange_sort_list.parameters())
optimizer2 = OPRO(strange_sort_list.parameters())
optimizer3 = TextGrad(strange_sort_list.parameters())
```

And you can easily implement your own optimizer that works directly with `TraceGraph` (more tutorials on how to work with TraceGraph coming soon).
Here is a summary of the optimizers:

| | Computation Graph | First-class Functions | Library Support | Speed | Large Graph |
|-------------------|-------------------|-----------------------|-----------------|----|-------------|
| OPRO |||| ⚡️ ||
| TextGrad |||| 🐌 ||
| OptoPrime (Trace) ||||||

The table evaluates the frameworks in the following aspects:

- Computation Graph: Whether the optimizer leverages the computation graph of the workflow.
- First-class Functions: Whether the framework allows users to write actual executable Python functions and not require
users to wrap them in strings.
- Library Support: Whether the framework has a library to support the optimizer.
- Speed: TextGrad is about 2-3x slower than OptoPrime (Trace). OPRO has no concept of computational graph, therefore is very fast.
- Large Graph: OptoPrime (Trace) represents the entire computation graph in context, therefore, will have issue with graphs that have more than hundreds of operations. TextGrad does not have such issue.

We provide a comparison to validate our implementation of TextGrad in Trace:

<p align="center">
<img src="https://github.com/microsoft/Trace/blob/main/docs/images/compare_to_textgrad.png" alt="drawing" width="90%"/>
</p>

You can also easily implement your own optimizer that works directly with `TraceGraph` (more tutorials on how to work
with TraceGraph coming soon).

## Citation

If you use this code in your research please cite the following [publication](https://arxiv.org/abs/2406.16218):

```
@article{cheng2024trace,
title={Trace is the New AutoDiff--Unlocking Efficient Optimization of Computational Workflows},
Expand All @@ -141,26 +182,37 @@ If you use this code in your research please cite the following [publication](ht
}
```


## Evaluation
A previous version of Trace was tested with gpt-4-0125-preview on numerical optimization, simulated traffic control, big-bench-hard, and llf-metaworld tasks, which demonstrated good optimization performance on multiple random seeds; please see the paper for details.

**Note** For gpt-4o, please use the version **gpt-4o-2024-08-06** (onwards), which [fixes](https://platform.openai.com/docs/models/gpt-4o) the structured output issue of gpt-4o-2024-05-13.
While gpt-4 works reliably most of the time, we've found gpt-4o-2024-05-13 often hallucinates even in very basic optimization problems and does not follow instructions. This might be due to the current implementation of optimizers rely on outputing in json format. Issues of gpt-4o with json have been reported in the communities (see [example](https://community.openai.com/t/gpt-4o-doesnt-consistently-respect-json-schema-on-tool-use/751125)).
A previous version of Trace was tested with gpt-4-0125-preview on numerical optimization, simulated traffic control,
big-bench-hard, and llf-metaworld tasks, which demonstrated good optimization performance on multiple random seeds;
please see the paper for details.

**Note** For gpt-4o, please use the version **gpt-4o-2024-08-06** (onwards),
which [fixes](https://platform.openai.com/docs/models/gpt-4o) the structured output issue of gpt-4o-2024-05-13.
While gpt-4 works reliably most of the time, we've found gpt-4o-2024-05-13 often hallucinates even in very basic
optimization problems and does not follow instructions. This might be due to the current implementation of optimizers
rely on outputing in json format. Issues of gpt-4o with json have been reported in the communities (
see [example](https://community.openai.com/t/gpt-4o-doesnt-consistently-respect-json-schema-on-tool-use/751125)).

## Disclaimers

- Trace is an LLM-based optimization framework for research purpose only.
- The current release is a beta version of the library. Features and more documentation will be added, and some functionalities may be changed in the future.
- System performance may vary by workflow, dataset, query, and response, and users are responsible for determining the accuracy of generated content.
- The current release is a beta version of the library. Features and more documentation will be added, and some
functionalities may be changed in the future.
- System performance may vary by workflow, dataset, query, and response, and users are responsible for determining the
accuracy of generated content.
- System outputs do not represent the opinions of Microsoft.
- All decisions leveraging outputs of the system should be made with human oversight and not be based solely on system outputs.
- Use of the system must comply with all applicable laws, regulations, and policies, including those pertaining to privacy and security.
- The system should not be used in highly regulated domains where inaccurate outputs could suggest actions that lead to injury or negatively impact an individual's legal, financial, or life opportunities.

- All decisions leveraging outputs of the system should be made with human oversight and not be based solely on system
outputs.
- Use of the system must comply with all applicable laws, regulations, and policies, including those pertaining to
privacy and security.
- The system should not be used in highly regulated domains where inaccurate outputs could suggest actions that lead to
injury or negatively impact an individual's legal, financial, or life opportunities.

## Contributing

This project welcomes contributions and suggestions. Most contributions require you to agree to a
This project welcomes contributions and suggestions. Most contributions require you to agree to a
Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us
the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.

Expand All @@ -177,10 +229,10 @@ contact [[email protected]](mailto:[email protected]) with any additio
This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft
trademarks or logos is subject to and must follow
[Microsoft's Trademark & Brand Guidelines](https://www.microsoft.com/en-us/legal/intellectualproperty/trademarks/usage/general).
Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship.
Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft
sponsorship.
Any use of third-party trademarks or logos are subject to those third-party's policies.


## Privacy

See [Microsoft Privacy Statement](https://privacy.microsoft.com/en-us/privacystatement).
Binary file added docs/images/compare_to_textgrad.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.

0 comments on commit 41be285

Please sign in to comment.