Trace is a Python library for end-to-end generative optimization of AI systems, enabling training with diverse feedback types like rewards, natural language, or compiler errors. It generalizes back-propagation by capturing and propagating execution traces, allowing users to optimize Python code directly, similar to training neural networks.

How It Works

Trace captures computation traces using node and bundle primitives. node wraps data that can be optimized, while @bundle decorates Python functions for optimization. The OptoPrime optimizer, along with others like TextGrad and OPRO, then updates parameters based on feedback, enabling optimization of complex AI agent behaviors and code generation.

Quick Start & Requirements

• Install: pip install trace-opt or pip install trace-opt[autogen] for AutoGen support.
• Requirements: Python >= 3.9. LLM API keys (OpenAI, Anthropic) are needed for LLM-based optimization. LiteLLM is the default backend; AutoGen can be used with an OAI_CONFIG_LIST file.
• Documentation: https://microsoft.github.io/Trace/
• Discord: https://discord.gg/h3b7b3h9

Highlighted Details

• Generalizes back-propagation for arbitrary feedback (rewards, text, errors).
• Supports multiple optimizers: OptoPrime (2-3x faster than TextGrad), TextGrad, and OPRO.
• Integrates with LLMs via LiteLLM or AutoGen for generative optimization tasks.
• Demonstrated applications include code optimization for parallel programming and AI agent training.

Maintenance & Community

• Actively developed by Microsoft.
• Featured in G-Research NeurIPS highlight; presented at NeurIPS 2024 Expo and UW Robotics Colloquium.
• Active Discord channel for support and discussion.
• Mailing list for announcements: https://groups.google.com/g/trace-opt-announcements

Licensing & Compatibility

• MIT License.
• Compatible with commercial use and closed-source linking.

Limitations & Caveats

Trace is a beta version for research purposes. LLM API performance can vary, and specific models (e.g., gpt-4o-2024-05-13) may exhibit issues with structured output. Large graphs (>hundreds of operations) might pose challenges for OptoPrime.