PaperlayerView exact revision on arXiv

Source-linked AI summary

TextGrad: Automatic "Differentiation" via Text

Mert Yuksekgonul, Federico Bianchi, Joseph Boen, Sheng Liu, Zhi Huang, Carlos Guestrin, James Zou

arXiv:2406.07496v1cs.CLcs.AIcs.LG

TL;DR

Compound AI systems built from interacting LLMs and tools require principled, automated optimization methods rather than expert-crafted heuristics. TextGrad uses LLM-generated textual feedback as gradients over computation graphs, improving performance across coding, question answering, reasoning, chemistry, and radiotherapy planning; in Google-Proof Question Answering, zero-shot GPT-4o accuracy rises from 51% to 55%.

  • Problem

    Compound AI systems increasingly combine multiple LLMs, tools, and complex components, but their optimization remains dependent on expert-crafted designs and heuristics.

  • Method

    TextGrad represents AI systems as computation graphs and propagates interpretable natural-language feedback from LLMs through arbitrary functions to optimize their variables.

  • Results

    TextGrad achieves improvements across coding, question answering, reasoning, molecule design, and radiotherapy planning, including a 4-percentage-point increase in zero-shot GPT-4o accuracy on Google-Proof Question Answering, from 51% to 55%.

  • Takeaways & Limitations

    TextGrad provides a general, PyTorch-like framework for optimizing diverse compound AI systems without being hand-crafted for a specific application domain.

  • Takeaways & Limitations

    Molecule and treatment-plan applications are validated only in silico; experimental and clinical assessments remain outside the paper’s scope.

Abstract

from arXiv · show

AI is undergoing a paradigm shift, with breakthroughs achieved by systems orchestrating multiple large language models (LLMs) and other complex components. As a result, developing principled and automated optimization methods for compound AI systems is one of the most important new challenges. Neural networks faced a similar challenge in its early days until backpropagation and automatic differentiation transformed the field by making optimization turn-key. Inspired by this, we introduce TextGrad, a powerful framework performing automatic ``differentiation'' via text. TextGrad backpropagates textual feedback provided by LLMs to improve individual components of a compound AI system. In our framework, LLMs provide rich, general, natural language suggestions to optimize variables in computation graphs, ranging from code snippets to molecular structures. TextGrad follows PyTorch's syntax and abstraction and is flexible and easy-to-use. It works out-of-the-box for a variety of tasks, where the users only provide the objective function without tuning components or prompts of the framework. We showcase TextGrad's effectiveness and generality across a diverse range of applications, from question answering and molecule optimization to radiotherapy treatment planning. Without modifying the framework, TextGrad improves the zero-shot accuracy of GPT-4o in Google-Proof Question Answering from $51\%$ to $55\%$, yields $20\%$ relative performance gain in optimizing LeetCode-Hard coding problem solutions, improves prompts for reasoning, designs new druglike small molecules with desirable in silico binding, and designs radiation oncology treatment plans with high specificity. TextGrad lays a foundation to accelerate the development of the next-generation of AI systems.

1 Introduction

Compound AI systems increasingly combine LLMs, tools, and other components, but their expert-crafted construction lacks principled automated optimization. TextGrad addresses this gap by propagating LLM-generated textual feedback through computation graphs and demonstrates improvements across diverse tasks.

  • Compound AI systems increasingly integrate multiple LLMs, tools, and sophisticated components.
  • Expert handcrafting and heuristic tuning remain common, motivating principled and automated optimization methods.
  • TextGrad treats textual LLM criticism as gradient-like feedback propagated through computation graphs containing arbitrary function calls.
  • The framework assumes current state-of-the-art LLMs can reason about individual components and subtasks, and applies across question answering, coding, reasoning, chemistry, and medicine.
  • 51% to 55%: TextGrad improves zero-shot GPT-4o accuracy on Google-Proof Question Answering while also producing gains in coding and other applications.

2 TEXTGRAD: Optimizing AI systems by backpropagating text feedback

TextGrad adapts automatic differentiation to compound systems by using natural-language feedback as gradients. Its domain-independent operators optimize variables in computation graphs, including unstructured data and nondifferentiable objectives.

  • Textual gradients: Textual gradients combine variable context with natural-language criticism describing how to improve the variable for a downstream objective.
  • Optimization: Textual Gradient Descent updates a variable using its current value and the feedback collected during the backward pass.
  • General framework: TextGrad represents compound AI systems as computation graphs whose variables are connected by LLMs, simulators, or numerical solvers.
  • Backpropagation: Feedback for a variable is collected from every successor context and aggregated recursively through the graph.
  • Objectives: The objective may be complex, nondifferentiable, and defined over unstructured data, with signals supplied by language-model evaluation, unit tests, or other tools.
  • Optimization settings: Instance optimization directly improves solutions, while both instance and prompt optimization can be performed without hand-crafting the framework.

3 Results

TextGrad is evaluated across code, question answering, prompt, molecule, and radiotherapy optimization, with improvements reported across these diverse tasks. The experiments use textual feedback to refine solutions, prompts, molecular structures, and treatment-plan objectives.

  • Code optimization: TextGrad optimizes code solutions at test time using LLM evaluation of the problem, code, instructions, and local test results.The code-refinement objective uses self-evaluation to critique each current iteration.
  • Prompt optimization: TextGrad significantly improves zero-shot prompts across three reasoning tasks, matches DSPy on Word Sorting and GSM8k, and exceeds DSPy by 7% on Object Counting.Combining DSPy demonstrations with the TextGrad instruction increased GSM8k accuracy to 82.1%.
  • Molecule optimization: Across 58 targets, TextGrad consistently generates molecules with improved binding affinity and druglikeness regardless of the starting fragment.Generated molecules show competitive affinity and druglikeness relative to clinically approved molecules while exhibiting unique structures.
  • Radiotherapy treatment plan optimization: TextGrad treatment plans achieved higher PTV mean dose and matching D95 while reducing mean doses to healthy organs relative to clinician-optimized plans.The optimization adjusts PTV and organ-at-risk weights to balance tumor targeting and organ protection.

4 Related work

Prior work automated prompt and language-model system optimization through programmatic frameworks, numerical gradients, and LLM-based critics or optimizers. TextGrad presents a broader framework that propagates natural-language feedback across varied optimization settings.

  • Prompt optimization has used few-shot selection, in-context learning, chain-of-thought, and ensembles to improve LLM performance.
  • White-box prompt-optimization methods use numerical gradients but require model-parameter access, limiting their use with closed-source models.
  • DSPy treats complex LLM-based systems as multilayer programs and optimizes them programmatically across question answering, reasoning, and prompt-optimization tasks.
  • TextGrad instead propagates natural-language feedback and unifies optimization of LLM chains, test-time training, self-refinement, momentum, batch optimization, and natural-language constraints.

5 Discussion

TextGrad is presented as a general, PyTorch-like framework that propagates LLM-generated textual feedback through computation graphs to optimize compound AI systems. The discussion reports broad application results while identifying extensions and validation needs for future work.

  • Discussion: The framework achieved state-of-the-art results in code optimization and PhD-level question answering, while also optimizing prompts and demonstrating scientific applications.The scientific applications include molecule development and treatment-plan optimization as proof-of-concept results.
  • Limitations and future work: Future work includes broader graph operations, optimization-stability techniques, and experimental or clinical validation of scientific applications.The paper specifically notes that molecule and treatment-plan results used in silico validation, whereas ultimate assessment lies outside its scope.
  • Discussion: TextGrad combines LLM reasoning with backpropagation-style decomposition to optimize compound AI systems.The framework is positioned as general, easy-to-use, and open-source rather than hand-crafted for one domain.
  • TextGrad framework: Variables receive natural-language gradients describing how they should change to improve downstream objectives.Gradients are propagated from successors to predecessors during a backward pass over the computation graph.
  • TextGrad framework: Role descriptions inject user knowledge and steer how variables are optimized.For example, defining a prediction as the final numerical answer guides updates toward numerical values.
  • TextGrad framework: TextGrad mirrors PyTorch-style autograd through composable operations with defined forward and backward computations.The backward procedure initializes gradients, topologically traverses the graph, and provides predecessor feedback based on downstream information.

A.4 Textual Gradient Descent Implementation

Textual Gradient Descent uses LLM feedback and contextual information to update text variables while supporting batch optimization, constraints, momentum, and in-context examples.

  • Optimizer: The optimizer receives feedback, variable roles, and context to improve prompts, solutions, code, or other text-based variables.Its prompt instructs the LLM to identify important and correct feedback even when feedback may be noisy.
  • Batch optimization: Batch optimization concatenates gradients from multiple instances so the optimizer sees feedback from multiple sources.The implementation aggregates per-example losses with tg.sum before calling backward().
  • Constrained optimization: Natural-language constraints guide optimizer updates, such as requiring a response to end with a specified answer format.Constraints are appended to the optimizer prompt through <CONSTRAINTS> tags and passed to TextualGradientDescent.
  • Momentum: Momentum preserves past variable iterations and appends them to the optimizer prompt.This adapts the standard momentum idea by providing previous textual values rather than numerically combining gradients.
  • In-context examples: In-context examples provide references for the properties that optimized variables should exhibit.The optimizer incorporates these examples when modifying a variable's role-described content.

C.1 Methodology

The LeetCode-Hard methodology evaluates TextGrad on difficult coding problems by repeatedly testing, critiquing, and updating generated solutions. The experiment averages results across five random seeds.

  • Baseline: Reflexion generates an initial solution, submits locally passing code to harder tests, and requests feedback when tests fail.The baseline was rerun from the authors' available code with minor execution-related changes and a re-extracted dataset.
  • Experimental setup: The evaluation uses 39 LeetCode-Hard coding problems and averages results from five runs with different seeds.Each optimization iteration uses one call for test-time loss evaluation, one for gradient collection, and one for code updates.
  • Task example: The coding task includes implementing minimumTime for a grid where cells become visitable only at their specified minimum times.The example asks for the minimum time to reach the bottom-right cell, or -1 if it cannot be reached.

D.1 Methodology

The question-answering methodology optimizes model solutions using structured multiple-choice prompts, answer-format constraints, and a critique-based objective. Evaluation extracts final answer letters and compares them with ground truth.

  • Prompt and constraints: The multiple-choice template requires step-by-step reasoning and a final line formatted as Answer: $LETTER.The allowed answer letters are A, B, C, and D.
  • Prompt and constraints: During optimization, the same final-answer constraint is supplied to the optimizer to preserve compatibility with evaluation.The constraint requires the response's last line to follow the specified answer format.
  • Evaluation: Evaluation uses string matching to extract one of A–D and compare it with the ground-truth answer.The GPQA Diamond, MMLU Machine Learning, and College Physics subsets contain 198, 112, and 92 questions, respectively.
  • Optimization objective: The objective asks a critical scientist to inspect model reasoning and identify why a prediction could be incorrect.It also prompts consideration of alternative answers through the Janusian Process.

E Prompt Optimization

The prompt-optimization appendix defines benchmark tasks, example queries, dataset splits, and task-specific evaluation procedures for optimizing reasoning prompts.

  • E.1 Tasks: Prompt optimization covers word sorting, object counting, and GSM8k-style arithmetic reasoning tasks.
  • E.1 Tasks: The word-sorting example asks a model to alphabetize a list containing nine words.
  • E.1 Tasks: The object-counting example asks how many objects appear in a household-item list containing two tables.
  • E.1 Tasks: The GSM8k example asks for Ahito’s miles after Amber ran 8 miles and Micah ran 3.5 times Amber’s distance.
  • E.1 Tasks: Word sorting and object counting use BBH data split into 50 training, 100 validation, and 100 test samples.
  • E.1 Tasks: GSM8k uses DSPy-provided splits containing 200 training, 300 validation, and 1319 test samples.
  • E.1 Tasks: Object counting and GSM8k use string-based exact match based on the last numerical answer value.
  • E.1 Tasks: A system prompt evaluates question-answering predictions by requiring only 1 or 0 inside <ACCURACY> tags.

F.1 Docking and Druglikeness Evaluation

Molecule evaluation converts generated SMILES into standardized chemical structures, then scores druglikeness and protein-binding affinity for textual optimization.

  • Generated SMILES are sanitized with RDKit before molecular properties are evaluated.
  • Valid molecules receive QED druglikeness scores from RDKit and Vina docking scores from the standardized DOCKSTRING workflow.
  • The molecule-optimization prompt asks an LLM to criticize a SMILES string using docking and druglikeness scores for a named protein target.
  • The prompt treats Vina scores below −10 as good and QED scores above 0.8 as good.
  • TEXTGRAD omits detailed binding-pocket geometry, retaining only target names and score-based feedback to preserve experimental generality.
  • The vina_qed_ratio is set to 10, making docking affinity the higher-priority objective than druglikeness.

F.3 Benchmarks

The benchmark section compares TEXTGRAD optimization across diverse starting fragments and targets, using post-selection to examine resulting molecule-score distributions.

  • Clinically approved comparator molecules come from DrugBank and are filtered for small-molecule, fully approved, orthosteric binders matching DOCKSTRING targets.
  • Supplementary Figure 1 initializes TEXTGRAD with fragments from three diverse functional groups and runs 10 optimization iterations across 58 DOCKSTRING targets.
  • For each fragment and protein, generated molecules undergo post-selection using the summary score before their distributions are visualized.
  • Although starting-fragment QED and Vina distributions vary greatly, optimized-molecule distributions are highly overlapping.
  • Large-scale pre-optimization screening is described as the practical drug-discovery setting in which promising leads are selected for further refinement.

F.5 Chemical Novelty

The chemical-novelty analysis measures similarity to known compounds and reports increasing novelty during optimization, alongside weak relationships between approved-drug similarity and performance.

  • A molecule is labeled novel when no ChEMBL compound has Tanimoto similarity greater than 0.8.
  • TEXTGRAD novelty increases over optimization updates, while Tversky similarity to clinically approved drugs varies widely across targets.
  • The analysis tests whether LLM-generated molecules might reproduce clinically approved structures from memorized training data.
  • Tversky similarity uses asymmetric substructure comparison, with DrugBank molecules as references and TEXTGRAD molecules as variants.
  • The selected analysis contains 87 generated molecules and 118 clinically approved drugs across 29 targets.
  • Tversky similarity has a median of 0.42 and ranges from 0.14 to 0.90.
  • Similarity to approved drugs is slightly anti-correlated with overall performance, while some generated molecules match or exceed comparator QED and Vina scores.
  • By iteration 6, 95% of molecules generated across 58 targets and three starting fragments are novel under the stated criterion.

F.6 Implicit Objectives

TEXTGRAD addresses safety and treatment-planning objectives through explicit evaluation and nested optimization. It evaluates generated molecules for harmfulness and optimizes radiotherapy plans against clinical protocols and dose objectives.

  • Molecule safety: TEXTGRAD molecules are evaluated for mutagenicity and clinical toxicity because these safety properties are not directly encoded in its objective.ADMET-AI predicts both properties from molecular structures.
  • Molecule safety: 87 generated molecules are compared with 118 clinically approved DrugBank molecules on predicted mutagenicity and clinical toxicity.The generated set contains the best-performing molecule for each protein target and initial fragment.
  • Radiotherapy treatment planning: Radiotherapy plan quality is assessed by an LLM that compares dose-volume histograms with clinical protocols and produces a final plan evaluation.The evaluation extracts relevant DVH entries, interprets them, checks protocol factors, and gives improvement suggestions when needed.
  • Radiotherapy treatment planning: TEXTGRAD optimizes importance weights in an outer loop while an inner numerical optimizer determines fluence maps by minimizing weighted target and organ-at-risk objectives.The cost function penalizes target-dose deviations and organ-at-risk overdosing, with positive fluence constraints and a D95 clinical constraint.
  • Radiotherapy treatment planning: The treatment-planning evaluation uses imaging and treatment-plan data from 5 prostate cancer patients who underwent IMRT.Each patient has CT scans, delineated anatomy, and a clinically approved plan.

G.4 Additional Results

Additional results compare TEXTGRAD-optimized treatment plans with clinician-optimized plans across target and organ-at-risk dose metrics. The reported tables summarize deviations and variability across 5 plans.

  • Comparison: The supplementary results directly compare TEXTGRAD-optimized plans with clinician-optimized plans.These comparisons provide additional treatment-planning results beyond the main evaluation.
  • PTV dose metrics: Supplementary Table 1 compares clinical and TEXTGRAD-optimized plans using PTV mean, minimum, and D95 dose metrics.It reports average deviations from clinical goals and standard deviations across 5 plans.
  • OAR dose metrics: Supplementary Table 2 reports mean-dose metrics for organs at risk to capture OAR sparing across 5 plans.Lower mean doses indicate better OAR sparing relative to clinical guidelines.
Loading 2406.07496v1…