CoChat
Sign Up
PreprintCited by 4Software Testing and Debugging TechniquesSoftware Engineering ResearchEvolutionary Algorithms and Applications

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

Open access
Cheng, Ching-An, Nie, Allen, Swaminathan, Adith·2024·arXiv
DOI 10.48550/arxiv.2406.16218arXiv 2406.16218
Open this paper in CoChat — free accountOpen PDF

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

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2406.16218

Cites

Statistics

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

  • TextGrad, Trace, and LLM-AutoDiff optimize arbitrary text components in chained workflows by back-propagating textual feedback using LLMs.

    “TextGrad [34], Trace [7], and LLM-AutoDiff [32] optimize arbitrary text components in chained workflows by back-propagating textual feedback using LLMs.”

    From HarnessOpt-Bench: Evaluating LLMs at Harness Optimization (2026) · cites this paper as [7]

    Section 5.5 explicitly distinguishes the methods: TextGrad propagates text feedback, whereas Trace propagates minimal subgraphs. Section 1.1 likewise says Trace propagates execution traces rather than gradients. Although Trace can optimize prompts in workflows, the mechanism attributed to Trace here is contrary to the paper. (LLM, conf=0.88)

    Checked by a reasoning model — may be wrong.

    Experimental — accuracy not yet measured

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Cites

No references could be resolved yet.

Research works printed but not linked: 45

  1. Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. arXiv preprint arXiv:2303.08774, 2023. · arXiv
  2. Albert Ziegler, Eirini Kalliamvakou, X. Alice Li, Andrew Rice, Devon Rifkin, Shawn Simis- ter, Ganesh Sittampalam, and Edward Aftandilian. Productivity assessment of neural code completion. In SIGPLAN International Symposium on Machine Programming, page 21-29, 2022.
  3. Jacky Liang, Wenlong Huang, Fei Xia, Peng Xu, Karol Hausman, Brian Ichter, Pete Florence, and Andy Zeng. Code as policies: Language model programs for embodied control. In IEEE International Conference on Robotics and Automation, pages 9493-9500, 2023.
  4. Qingyun Wu, Gagan Bansal, Jieyu Zhang, Yiran Wu, Shaokun Zhang, Erkang Zhu, Beibin Li, Li Jiang, Xiaoyun Zhang, and Chi Wang. Autogen: Enabling next-gen LLM applications via multi-agent conversation framework. arXiv preprint arXiv:2308.08155, 2023. · arXiv
  5. Frank Hutter, Holger H Hoos, and Kevin Leyton-Brown. Sequential model-based optimization for general algorithm configuration. In Learning and Intelligent Optimization, pages 507-523. Springer, 2011.
  6. Peter I Frazier. Bayesian optimization. Recent Advances in Optimization and Modeling of Contemporary Problems, pages 255-278, 2018.
  7. Thomas Bäck and Hans-Paul Schwefel. An overview of evolutionary algorithms for parameter optimization. Evolutionary computation, 1(1):1-23, 1993.
  8. Josep Ginebra and Murray K Clayton. Response surface bandits. Journal of the Royal Statistical Society Series B: Statistical Methodology, 57(4):771-784, 1995.
  9. Yongchao Zhou, Andrei Ioan Muresanu, Ziwen Han, Keiran Paster, Silviu Pitis, Harris Chan, and Jimmy Ba. Large language models are human-level prompt engineers. In ICLR, 2023.
  10. Reid Pryzant, Dan Iter, Jerry Li, Yin Lee, Chenguang Zhu, and Michael Zeng. Automatic prompt optimization with "gradient descent" and beam search. In EMNLP, pages 7957-7968, 2023.
  11. Chengrun Yang, Xuezhi Wang, Yifeng Lu, Hanxiao Liu, Quoc V Le, Denny Zhou, and Xinyun Chen. Large language models as optimizers. In ICLR, 2024.
  12. Tobias Schnabel and Jennifer Neville. Prompts as programs: A structure-aware approach to efficient compile-time prompt optimization. arXiv preprint arXiv:2404.02319, 2024. · arXiv
  13. Eric Zelikman, Eliana Lorch, Lester Mackey, and Adam Tauman Kalai. Self-taught optimizer (STOP): Recursively self-improving code generation. arXiv preprint arXiv:2310.02304, 2023. · arXiv
  14. Allen Nie, Ching-An Cheng, Andrey Kolobov, and Adith Swaminathan. The importance of directional feedback for LLM-based optimizers. arXiv preprint arXiv:2405.16434, 2024. · arXiv
  15. David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-propagating errors. Nature, 323(6088):533-536, 1986.
  16. Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. Journal of Machine Learning Research, 18(153):1-43, 2018.
  17. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: an imperative style, high-performance deep learning library. In NeurIPS, pages 8026-8037, 2019.
  18. Karthik Nagaraj, Charles Killian, and Jennifer Neville. Structured comparative analysis of systems logs to diagnose performance problems. In USENIX Symposium on Networked Systems Design and Implementation, pages 353-366, 2012.
  19. Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc Le, Ed Chi, Denny Zhou, et al. Challenging BIG-Bench tasks and whether chain-of-thought can solve them. In Findings of the ACL, pages 13003-13051, 2023.
  20. Omar Khattab, Arnav Singhvi, Paridhi Maheshwari, Zhiyuan Zhang, Keshav Santhanam, Sri Vardhamanan, Saiful Haq, Ashutosh Sharma, Thomas T Joshi, Hanna Moazam, et al. DSPy: Compiling declarative language model calls into self-improving pipelines. arXiv preprint arXiv:2310.03714, 2023. · arXiv
  21. Ching-An Cheng, Andrey Kolobov, Dipendra Misra, Allen Nie, and Adith Swaminathan. LLF-Bench: Benchmark for interactive learning from language feedback. arXiv preprint arXiv:2312.06853, 2023. · arXiv
  22. Mert Yuksekgonul, Federico Bianchi, Joseph Boen, Sheng Liu, Zhi Huang, Carlos Guestrin, and James Zou. TextGrad: Automatic "differentiation" via text. arXiv preprint arXiv:2406.07496, 2024. · arXiv
  23. Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. · arXiv
  24. Toru Seo. UXsim: An open source macroscopic and mesoscopic traffic simulator in python-a technical overview. arXiv preprint arXiv:2309.17114, 2023. · arXiv
  25. Courtney Slavin, Wei Feng, Miguel Figliozzi, and Peter Koonce. Statistical study of the impact of adaptive traffic signal control on traffic and transit performance. Transportation Research Record, 2356(1):117-126, 2013.
  26. James Kennedy and Russell Eberhart. Particle swarm optimization. In ICNN, volume 4, pages 1942-1948, 1995.
  27. Xiao Liu, Hao Yu, Hanchen Zhang, Yifan Xu, Xuanyu Lei, Hanyu Lai, Yu Gu, Hangliang Ding, Kaiwen Men, Kejuan Yang, Shudan Zhang, Xiang Deng, Aohan Zeng, Zhengxiao Du, Chenhui Zhang, Sheng Shen, Tianjun Zhang, Yu Su, Huan Sun, Minlie Huang, Yuxiao Dong, and Jie Tang. AgentBench: Evaluating LLMs as agents. In ICLR, 2024.
  28. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H Chi, Quoc V Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. In NeurIPS, pages 24824-24837, 2022.
  29. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In NeurIPS, 2020.
  30. Timo Schick, Jane Dwivedi-Yu, Roberto Dessi, Roberta Raileanu, Maria Lomeli, Eric Hambro, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools. In NeurIPS, 2023.
  31. Weiran Yao, Shelby Heinecke, Juan Carlos Niebles, Zhiwei Liu, Yihao Feng, Le Xue, Rithesh R N, Zeyuan Chen, Jianguo Zhang, Devansh Arpit, Ran Xu, Phil L Mui, Huan Wang, Caiming Xiong, and Silvio Savarese. Retroformer: Retrospective large language agents with policy gradient optimization. In ICLR, 2024.
  32. Alessandro Sordoni, Xingdi Yuan, Marc-Alexandre Côté, Matheus Pereira, Adam Trischler, Ziang Xiao, Arian Hosseini, Friederike Niedtner, and Nicolas Le Roux. Joint prompt optimiza- tion of stacked LLMs using variational inference. In NeurIPS, 2023.
  33. Mingchen Zhuge, Wenyi Wang, Louis Kirsch, Francesco Faccio, Dmitrii Khizbullin, and Jurgen Schmidhuber. Language agents as optimizable graphs. arXiv preprint arXiv:2402.16823, 2024. · arXiv
  34. Ruotian Ma, Xiaolei Wang, Xin Zhou, Jian Li, Nan Du, Tao Gui, Qi Zhang, and Xuanjing Huang. Are large language models good prompt optimizers? arXiv preprint arXiv:2402.02101, 2024. · arXiv
  35. Xinyuan Wang, Chenxi Li, Zhen Wang, Fan Bai, Haotian Luo, Jiayou Zhang, Nebojsa Jojic, Eric Xing, and Zhiting Hu. PromptAgent: Strategic planning with language models enables expert-level prompt optimization. In ICLR, 2024.
  36. Xinyu Tang, Xiaolei Wang, Wayne Xin Zhao, Siyuan Lu, Yaliang Li, and Ji-Rong Wen. Unleashing the potential of large language models as prompt optimizers: An analogical analysis with gradient-based model optimizers. arXiv preprint arXiv:2402.17564, 2024. · arXiv
  37. Xinyun Chen, Maxwell Lin, Nathanael Schärli, and Denny Zhou. Teaching large language models to self-debug. In ICLR, 2024.
  38. Zeyuan Ma, Hongshu Guo, Jiacheng Chen, Guojun Peng, Zhiguang Cao, Yining Ma, and Yue-Jiao Gong. LLaMoCo: Instruction tuning of large language models for optimization code generation. arXiv preprint arXiv:2403.01131, 2024. · arXiv
  39. Tennison Liu, Nicolás Astorga, Nabeel Seedat, and Mihaela van der Schaar. Large language models to enhance bayesian optimization. In ICLR, 2024.
  40. Michael R. Zhang, Nishkrit Desai, Juhan Bae, Jonathan Lorraine, and Jimmy Ba. Using large language models for hyperparameter optimization. arXiv preprint arXiv:2312.04528, 2023. · arXiv
  41. Gábor Bartók, Dean P Foster, Dávid Pál, Alexander Rakhlin, and Csaba Szepesvári. Partial monitoring-classification, regret bounds, and algorithms. Mathematics of Operations Research, pages 967-997, 2014.
  42. Laurent Hascoet and Mauricio Araya-Polo. Enabling user-driven checkpointing strategies in reverse-mode automatic differentiation. arXiv preprint cs/0606042, 2006. · arXiv
  43. Amirreza Shaban, Ching-An Cheng, Nathan Hatch, and Byron Boots. Truncated back- propagation for bilevel optimization. In AISTATS, pages 1723-1732, 2019.
  44. Ching-An Cheng, Jonathan Lee, Ken Goldberg, and Byron Boots. Online learning with continuous variations: Dynamic regret and reductions. In AISTATS, pages 2218-2228, 2020.
  45. Sham Machandranath Kakade. On the Sample Complexity of Reinforcement Learning. PhD thesis, University College London, 2003.

Cited by

Sorted by citations

Authors
Cheng, C.-A., Nie, A., & Swaminathan, A.
ORCID
Venue
arXiv
Publisher
arXiv
Published
2024
Pages
71596-71642
ISSN
1049-5258 (ISSN-L)
DOI
10.48550/arxiv.2406.16218
arXiv
2406.16218
OpenAlex
W4415797156
Cited by
4
References
0
Citations by year
2026: 2 · 2025: 2
FWCI
1.09
Type
conference paper
Language
en
Open access
Licence
Locations
  • neural information processing systems (publishedVersion)
  • arXiv (Cornell University) (submittedVersion, cc-by)
  • arXiv (Cornell University)
Funders
Microsoft Research
Topics
  • Physical Sciences › Computer Science › Software › Software Testing and Debugging Techniques
  • Physical Sciences › Computer Science › Information Systems › Software Engineering Research
  • Physical Sciences › Computer Science › Artificial Intelligence › Evolutionary Algorithms and Applications
Keywords
traces, automatic differentiation, execution traces, large language models, program optimization, feedback-driven optimization

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review