斯坦福大学最新大模型ChatGPT提示优化框架TextGrad

TEXTGRAD：用自然语言反馈“反向传播”来优化复合AI系统

AI领域正在经历一场范式转变——突破不再来自单一模型，而是来自多个大语言模型（LLMs）与其他复杂组件的编排组合。因此，如何为这类复合AI系统设计合理且自动的优化方法，成了当前最紧迫的挑战之一。

回顾历史，神经网络早期也面临过类似困境，直到反向传播和自动微分出现，才让优化成为核心驱动力，彻底改变了整个领域。受此启发，我们来看一项新工作：

——一个通过文本执行自动“微分”的框架。它利用LLMs提供的文本反馈进行反向传播，从而逐步改进复合AI系统中各个组件。在这个框架里，LLMs可以为计算图中的变量（从代码片段到分子结构）生成丰富、通用、且用自然语言表达的优化建议。

TEXTGRAD遵循PyTorch的语法和抽象，灵活且易用。只要用户提供目标函数，无需调整框架组件或提示，就能开箱即用。我们在多个任务上验证了它的有效性：从问答、分子优化到放射治疗计划。例如，在不修改框架的前提下，TEXTGRAD将GPT-4o在Google-Proof问题回答中的零样本准确率从51%提升到55%，在LeetCode-Hard编码问题解决方案上获得20%的相对性能提升，改进了推理提示，设计了具有理想体内结合力的类药小分子，并做出了高特异性的放射肿瘤治疗计划。TEXTGRAD为加速下一代AI系统的开发奠定了基础。

原文：Automatic "Differentiation" via Text
地址：https://arxiv.org/abs/2406.07496
代码：https://github.com/zou-group/textgrad
出版：arXiv
机构: 斯坦福大学

1 研究问题

本文的核心问题是：如何设计一个通用且自动的框架，来优化由多个大语言模型及其他复杂组件构成的复合AI系统。

举个具体的例子：一个由两部分组成的问答系统——第一部分是LLM，根据给定提示（Prompt）和问题（Question）生成答案（Prediction）；第二部分是另一个LLM，根据评估指令（Evaluation Instruction）和第一部分的答案给出评分（Evaluation）。现在的挑战是：如何自动调整第一部分的提示，让整个系统的最终评分尽可能高？

这个问题之所以棘手，主要源于几个特性：

• 复合AI系统中的组件（如LLM）通常是不可微的黑盒子，无法直接计算梯度。
• 中间变量和输出往往是非结构化的，比如自然语言文本。
• 目标函数可能本身就很复杂、不可微，甚至需要另一个AI系统来计算。
• 针对不同场景，现有方法往往需要大量定制开发，缺乏一个通用的、具备“自动微分”能力的优化框架。

针对这些挑战，TEXTGRAD提出了一种基于“文本反馈自动微分”的思路。

它的核心是将复合AI系统表示为一个计算图，图中的变量可以是任意非结构化数据（代码、分子式、自然语言文本等），边上的函数可以是任意复杂的不可微算子（如LLM调用、数值模拟器等）。给定一个下游目标后，TEXTGRAD会通过LLM产生的文本反馈（称为“文本梯度”）来自动指导每个变量的优化方向。就像PyTorch等框架通过链式法则自动计算导数一样，TEXTGRAD也能利用文本梯度的“链式传播”实现端到端的优化。对于图中的每条边（每个子函数），框架会提示LLM分析其输入、输出以及已有的输出梯度，进而“反向传播”出输入变量的梯度。通过迭代优化，最终让整个系统在目标函数上取得进展。值得一提的是，这个框架可以无缝集成外部工具（如代码解释器、分子对接模拟器），并在多种场景下开箱即用。

2 研究方法

2.1 TEXTGRAD框架综述

TEXTGRAD通过反向传播自然语言反馈来优化AI系统。它把系统表示成一个计算图，节点是变量，边是计算。变量可以是输入、输出或中间结果，值可以是文本、图像等非结构化数据，函数可以是LLM调用、仿真器或其他工具。

举个栗子：一个由两个LLM调用组成的简单系统：

其中第一个LLM根据提示和问题生成预测，第二个LLM根据指令对预测进行评估。我们希望优化提示来改善整体评估分数。

TEXTGRAD的核心是为变量计算自然语言“梯度”，这个梯度就是告诉我们需要如何修改变量值才能改善目标。反馈由LLM生成，形式比如“这个预测可以通过……来改进”。

得到梯度后，TEXTGRAD使用Textual Gradient Descent（TGD）来更新变量。TGD基于LLM，以当前变量值和梯度反馈为输入，生成新的变量值。直观上，LLM充当了“优化器”，以智能的方式修改变量。

总的来说，TEXTGRAD通过LLM提供的自然语言反馈，实现了类似于PyTorch中自动微分和优化的能力，从而灵活优化各种复杂AI系统。

2.2 TEXTGRAD的一般形式

前面以两个LLM调用的例子说明，但实际上它可以应用于任意复杂的计算图。对于图中任意变量，它的梯度由其所有后继节点的梯度经过该函数的backward函数计算得到。接受当前变量的值、所有后继节点的值以及它们的梯度作为输入，生成反馈。要计算一个变量的反馈，需要看它在图中所有直接用途，也就是如何参与生成后继，以及还需要如何改进。

得到任意变量的梯度后，就可以用TGD更新。通过前向传播计算变量值，反向传播计算梯度，并用TGD更新，整个系统就能被优化。这很像PyTorch等框架训练神经网络的过程。

2.3 实例优化与提示优化

TEXTGRAD主要关注两类优化问题：实例优化和提示优化。

实例优化中，将问题的解（如代码、答案、分子结构）作为变量，直接优化这个解以改进目标。例如，给定一个编程问题，把代码解构建为变量，通过TEXTGRAD迭代优化代码，使其在测试用例上表现更好。

提示优化中，则希望找到一个LLM的输入提示，使其在某个任务上整体表现更好。例如，找到一个数学题的提示，让LLM在所有数学问题上都表现好，而不是某个特定问题。这里优化变量是提示，loss在多个样本上评估。通过优化，得到的提示可以很好地泛化。

尽管目标不同，但在TEXTGRAD中，处理方式是统一的。实例优化可看作只有一个样本的提示优化。

2.4 优化技术

除了基础TGD更新，TEXTGRAD还支持多种优化技术。比如batch optimization，一次评估多个样本的loss，综合起来传播梯度（通过tg.sum操作实现，类似PyTorch的torch.sum）。

还有constrained optimization，通过自然语言描述约束条件，在变量更新时让LLM满足这些条件，比如要求“答案最后一行必须以‘Answer: $VALUE’格式给出”。这能更好地控制生成结果。

TEXTGRAD还支持momentum，更新变量时参考之前的变量值，平滑优化轨迹。这些技术大多借鉴了经典梯度下降法的思想，只不过用自然语言反馈替代了数值梯度。

通过以上设计，TEXTGRAD成为一个通用的优化框架，为构建下一代AI系统奠定了基础。

3 实验

3.1 实验场景介绍

实验在多个场景下展示TEXTGRAD的灵活性和有效性，包括代码优化、科学问题解答、推理任务、药物分子设计和放射治疗计划优化。复杂任务被表示为计算图，利用LLM的自然语言反馈进行优化，无需针对特定任务手工调整框架。

3.2 实验设置

• Datasets：LeetCode Hard、GPQA、MMLU、类药性数据库、前列腺癌患者CT影像
• Baseline：Reflexion、零样本CoT、DSPy等
• Implementation details：使用gpt-4o、gpt-3.5作为前向和后向传播模型，设计自然语言loss函数，迭代优化变量
• Metric：通过率、准确率、相似性、类药性、剂量分布等

3.3 实验结果

3.3.1 实验一、代码优化

目的：在LeetCode Hard数据集上优化代码，提高gpt-4o的性能。
涉及图表：表1
实验细节：以代码为优化变量，利用局部测试信息和自我评估迭代优化，每次迭代3次LLM调用。
结果：TEXTGRAD将gpt-4o的通过率从23%提高到36%，超过Reflexion方法的31%。

3.3.2 实验二、科学问题解答优化

目的：在科学问题数据集上优化解答，提高gpt-4o的zero-shot表现。
涉及图表：表2
实验细节：以答案为优化变量，通过LLM评估迭代优化，每次迭代3次，最后多数投票得到最终答案。
结果：在GPQA上将gpt-4o准确率从51%提高到55%，在MMLU子集上也有2-4个点的提升。

3.3.3 实验三、推理任务提示优化

目的：通过优化提示改进gpt-3.5在推理任务上的性能。
涉及图表：表3
实验细节：以系统提示为优化变量，用gpt-4o评估反馈，批量优化，共36个训练样本，最后在验证集挑选最优提示。
结果：在3个任务上均显著优于零样本CoT提示，与DSPy方法相当或更优。

3.3.4 实验四、药物分子优化

目的：同时优化分子对靶点的结合能和类药性。
涉及图表：图2、补充图1-3
实验细节：以SMILES分子式为优化变量，用Autodock Vina和QED评分计算loss，迭代10次优化3种初始分子片段。
结果：生成分子在58个靶点上持续改进结合能和类药性，在29个靶点上与临床药物相当，同时具有结构新颖性和安全性。

3.3.5 实验五、放射治疗计划优化

目的：优化前列腺癌放射治疗中的剂量分布，兼顾肿瘤剂量和保护器官。
涉及图表：图3、补充表1-2
实验细节：以PTV和OAR的权重超参为优化变量，内层用数值优化器matRad得到计划，外层用LLM计算loss，迭代优化权重。
结果：TEXTGRAD对PTV的剂量更加符合临床目标，对直肠膀胱等器官的保护优于人工计划。

4 总结后记

本文针对优化包含多个LLM和其他复杂组件的复合AI系统，提出了一种基于文本“反向传播”的自动优化框架TextGrad。它让LLM为计算图中的变量提供富有洞见的自然语言反馈，通过反向传播这些“文本梯度”来优化系统组件。TextGrad采用与PyTorch类似的语法和抽象，使用灵活方便。多个领域的应用表明，无需调整框架即可提升GPT-4等模型的零样本性能，为加速下一代AI系统的开发奠定了基础。

一些思考和延伸方向：

1. 除了使用LLM提供反馈，是否可以将TextGrad与其他形式的Critics（如判别器网络）结合，进一步提升优化性能？
2. 当前工作主要聚焦于实例优化，如何将TextGrad扩展到参数优化，实现端到端的可微分训练？
3. 计算图中包含的组件尚不够丰富，如何纳入更多类型的工具（如代码解释器、物理模拟器），增强框架适用范围？

可借鉴的方法点：

1. 以自动微分为类比，为优化LLM构成的复杂系统提供了新思路，为设计下一代优化算法开辟了空间。
2. 利用LLM提供带有推理和洞见的自然语言反馈，可用于指导其他机器学习系统的训练和优化。
3. 通过定义合适的抽象和接口，使优化过程可复用、模块化，降低开发门槛，这种范式可推广到其他AI开发场景。