【摘要】AI正在经历一场范式转变,突破是通过编排多个大语言模型(LLMs)和其他复杂组件的系统实现的。因此,开发复合AI系统的合理自动优化方法是最重要的新挑战之一。神经网络在其早期也面临类似的挑战,直到反向传播和自动微分通过使优化成为关键而改变了该领域。受此启发,我们引入了TEXTGRAD,一个通过文本执行自动"微分"的强大框架。TEXTGRAD通过LLMs提供的文本反馈来反向传播,以改进复合AI系统的各个组件。在我们的框架中,LLMs为计算图中的变量(从代码片段到分子结构)提供丰富的、通用的、自然语言的优化建议。TEXTGRAD遵循PyTorch的语法和抽象,灵活易用。它可以开箱即用地应用于各种任务,用户只需提供目标函数,而无需调整框架的组件或提示。我们展示了TEXTGRAD在从问题回答和分子优化到放射治疗计划的各种应用中的有效性和通用性。在不修改框架的情况下,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系统。

本文研究问题的特点和现有方法面临的挑战主要体现在以下几个方面:

• 复合AI系统的组件(如大语言模型)通常是不可微的黑盒子,无法直接计算梯度。这使得传统的基于梯度的优化方法难以直接应用。
• 复合AI系统的中间变量和输出通常是非结构化的,如自然语言文本。这种数据格式给优化过程带来额外的复杂性。
• 优化的目标函数可能是复杂的、不可微的,甚至需要另一个AI系统(如语言模型)来计算。这进一步增加了问题的难度。
• 针对不同的应用场景,现有方法通常需要大量的定制开发。缺乏一个通用的、具备"自动微分"能力的优化框架供研究人员使用。

针对这些挑战,本文提出了一种灵活而强大的"基于文本反馈的自动微分"框架TEXTGRAD:

2 研究方法

2.1 TEXTGRAD框架综述

本文提出了一种称为TEXTGRAD的自动优化框架,通过反向传播自然语言反馈来优化复杂的AI系统。TEXTGRAD将AI系统表示为一个计算图 ,其中节点 表示变量,边表示计算 。这里的变量可以是输入、输出或中间结果,值可以是文本、图像等非结构化数据,函数 可以是LLM调用、仿真器或其他工具。

举个糖炒栗子,考虑一个简单的由两个LLM调用 组成的系统:

其中 根据提示(Prompt)和问题生成预测,而 根据指令对预测进行评估。我们希望通过优化提示来改善整体评估分数。

TEXTGRAD的核心是为变量 计算自然语言"梯度" ,其中 是我们希望优化的下游目标。这里的"梯度"是指对变量的反馈,告诉我们如何修改 的值才能改善 。这个反馈由LLM生成,形式为自然语句,比如"这个预测可以通过......来改进"。

得到梯度反馈后,TEXTGRAD使用Textual Gradient Descent (TGD)来更新变量:

TGD基于LLM,以当前变量值 和梯度反馈 为输入,生成新的变量值 。直觉上,这就像是LLM充当了"优化器",以智能的方式修改。

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

2.2 TEXTGRAD的一般形式

前面我们以一个简单的两个LLM调用的例子说明了TEXTGRAD,但它可以应用于任意复杂的计算图。

具体来说,对于图中任意变量 ,假设其前驱节点(直接输入)为 ,后继节点(直接输出)为 ,即有:

则 的梯度由其所有后继节点的梯度经过 的backward函数 计算得到:

接受 的值、所有后继节点 的值、以及它们的梯度 为输入,生成 的梯度反馈。直觉上,要计算 的反馈,我们需要看 在图中的所有直接用途,即"如何参与生成 ",以及"还需要如何改进"。综合来看 还需要如何修改。

在得到任意变量的梯度之后,就可以用TGD进行更新:

这样,通过执行前向传播(forward)计算变量值,反向传播(backward)计算梯度,并用TGD更新变量,TEXTGRAD就可以优化整个AI系统。这非常类似于基于PyTorch等框架训练神经网络的过程。

2.3 实例优化与提示优化

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

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

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

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

2.4 优化技术

除了基础的TGD更新外,TEXTGRAD还支持多种优化技术,以提升效率和性能。

首先是batch optimization,即一次评估多个样本的loss,并将其综合起来传播梯度。这可以通过TEXTGRAD中的tg.sum操作实现,类似于PyTorch中的torch.sum。

其次是constrained optimization,通过自然语言描述约束条件,在变量更新时让LLM满足这些条件。例如,我们可以要求"答案的最后一行必须以'Answer: $VALUE'的格式给出,其中VALUE是一个数值"。加入约束可以更好地控制生成结果。

TEXTGRAD还支持momentum,在更新变量时参考之前的变量值。例如TGD可以接受之前步的值作为输入,以平滑优化轨迹。

这些优化技术大多借鉴了经典的梯度下降法的思想。在TEXTGRAD中,我们将其用自然语言反馈替代了数值梯度,但核心思想是一致的。

通过以上各部分,TEXTGRAD实现了一个非常通用的框架,可以优化各种复杂的AI系统。这为构建下一代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类似的语法和抽象,使用灵活方便。在编程、问答、分子优化、放射治疗计划等多个领域的应用表明,无需调整框架,TextGrad即可提升GPT-4等语言模型的零样本性能,为加速下一代AI系统的开发奠定了基础。

疑惑和想法:

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

可借鉴的方法点:

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