人工智能咨询培训老师叶梓 转载标明出处

在自动化设计和更新AI系统，如编码助手、机器人和副驾驶等领域，优化问题的研究显得尤为重要。传统的自动微分框架，例如PyTorch，虽然能够高效地优化可微分系统，但在面对一般计算工作流时，它们往往无法处理非可微分和涉及丰富反馈（例如控制台输出或用户响应）的情况。针对这一挑战，微软研究院和斯坦福大学的研究人员提出了一种新的优化框架——Trace，它能够利用大型语言模型（LLMs）等生成模型，对一般计算工作流进行端到端的生成优化。

想要掌握如何将大模型的力量发挥到极致吗？叶老师带您深入了解 Llama Factory —— 一款革命性的大模型微调工具（限时免费）。

1小时实战课程，您将学习到如何轻松上手并有效利用 Llama Factory 来微调您的模型，以发挥其最大潜力。

CSDN教学平台录播地址：https://edu.csdn.net/course/detail/39987

方法

Trace的框架基于一个新的数学模型——OPTO（Optimization with Trace Oracle），旨在解决涉及大模型的计算工作流优化问题。这一方法的核心在于将计算工作流的优化问题转化为OPTO问题，并利用Trace框架实现OPTO的迭代优化过程。

OPTO是Trace框架的基础，它定义了一种基于图的迭代优化抽象。这里计算图的概念被重新审视，如图2c所示，计算图是一个有向无环图（DAG），其中节点代表对象（如张量、字符串等），边表示输入输出关系。在优化的背景下，一些输入被标记为可训练参数，记作{Xθ}。对于节点X，其父节点是创建X的操作的输入，而X的后代是通过跟随有向边从X可达的节点；祖先节点则相反定义。为了简化，假设所有计算操作都有单一输出，这样可以不显式表示操作，而将操作与子节点关联起来，将完整的计算紧凑地表示为DAG。

OPTO问题实例由三元组（Θ, ω, T）定义，其中Θ是参数空间，ω是问题的上下文，T是Trace Oracle。在每次迭代中，优化器选择一个参数θ∈Θ，然后Trace Oracle T返回一个迹反馈τ = (f, g)，其中g是以DAG形式表示的执行迹线（参数包含在g的根节点中），f是对g的输出节点之一提供的反馈。优化器使用迹反馈τ根据上下文ω更新参数，并进入下一次迭代，如图3所示。在OPTO中，输出反馈f是通用的，可以是分数、梯度、用自然语言表达的提示/解释和控制台消息。上下文ω提供了解释输出反馈f的不变信息以及任何已知的边信息，例如参数的期望属性。上下文ω对于一个OPTO问题实例是固定的，而输出反馈f可以随着参数θ∈Θ和计算g的变化而变化。

Trace框架基于两个原语设计：node和bundle。node是Python对象的包装器，当包装时，Python对象在Trace的计算图中注册为一个唯一的节点。一个节点可以被设置为可训练的，使其成为OPTO中的一个参数。bundle是一个装饰器，将Python方法转换为操作。当一个函数被装饰时，它的文档字符串和源代码被记录为操作的定义，这推断了输出反馈应如何改变参数。

使用Trace涉及以下步骤：首先，用户使用node和bundle声明工作流的参数，并使用bundle定义工作流的概念块作为操作。然后，用户创建一个OPTO优化器，并可选地为问题提供上下文ω。用户还定义了一种机制，以提供对计算结果的反馈，例如分数、自然语言建议等。在声明之后，通过Trace的优化重复以下步骤：1）执行被装饰的工作流。在运行时，后端构建了一个DAG，记录计算结果及其连接。2）通过调用backward启动输出反馈到参数的传播。在内部，Trace提取连接参数和输出的最小子图g，并将OPTO优化器的迹反馈τ = (f, g)发送出去。3）调用OPTO优化器的step方法更新参数。

本文介绍了一个基于大模型的生成优化算法OptoPrime，用于任何基于文本的OPTO问题。OptoPrime的核心挑战之一是如何将执行迹线子图g（可能涉及各种图结构和异构数据）表示给大模型，以便大模型能够理解和推理参数更新的下游效应。利用大模型在编码和调试方面的卓越能力。OptoPrime通过Trace计算的迹反馈（f, g）作为一个伪算法问题来呈现：子图g被表示为代码执行报告，包含有关计算值和g中操作的描述的信息。然后提示大模型根据给定的输出反馈f更新g中的参数。

图4显示了一个例子。需要注意的是，尽管这些行看起来像实际程序，但它们不是实际程序，而是由bundle定义的计算图。参数更新方面，使用ReActCoT风格的提示在一次查询中提示大模型，要求其基于图生成推理，并提出参数更改建议。如果可以从响应中提取建议，将更新参数。为了稳健性，OptoPrime有一个基本的记忆模块，跟踪过去的参数-反馈对作为上下文示例。

实验

实验部分首先通过与现有的最先进大模型优化器OPRO进行比较，来验证OptoPrime的能力。OPRO不使用执行迹线，而是依赖于参数和反馈对的记忆。实验中使用的是GPT-4-0125-Preview模型，所有实验都在配备16GB RAM的标准个人电脑上进行，Trace在执行工作流时没有引入可测量的额外开销。

实验的第一部分是验证OptoPrime是否能够解决传统的可微分优化问题，因为这些问题是OPTO的一个特例。实验构建了一个合成任务环境，随机生成目标值y和具有任意复杂数值变量连接的计算图h。

实验评估了OptoPrime（记作Trace）和一个不可见图的变体（Trace Masked），这两个优化器都不使用记忆。输出反馈是“输出应该更大/更小”，这与相对于h的梯度具有相同的信息。实验还比较了Adam优化器的性能。在30次不同随机生成的问题上运行这些方法，所有方法都面对相同的随机性。结果表明，Trace能够与一流的Adam优化器相匹配； 另一方面，没有完整计算图的支持，仅凭反馈的优化器很难找到最优解x（图5a）。

实验的第二部分测试了Trace在交通控制问题中的性能，这是一个超参数调整的实例。实验使用UXSim模拟了一个四路交叉口的交通情况，其中可训练参数是两个介于[15, 90]之间的整数，代表每个方向交通流的绿灯持续时间。反馈是所有车辆因交叉口而经历的估计延迟，优化器的目标是使用最少的交通模拟次数来最小化延迟。为此，优化器必须找到时间分布和变化需求之间的正确权衡。

图5展示了交通控制文献中的SOTA启发式方法SCATS以及两种黑盒优化技术：高斯过程最小化（GP）和粒子群优化（PSO）的性能。所有方法都使用相同的起始参数。Trace表示使用记忆的OptoPrime，而Trace NoMem表示不使用记忆的OptoPrime。

实验结果表明，Trace很快就能与SCATS启发式方法相竞争，而OPRO则不能。 此外，实验发现记忆对于Trace在这个任务中的表现至关重要。但值得注意的是，与其他方法相比，Trace消耗了额外的开销，因为Trace需要具体化结果计算图，并用比OPRO更长的提示查询大模型。

许多当今的大模型代理，例如LangChain和DSPy指定的代理，都有许多组件。这些库提供了优化工具来优化工作流的一小部分，主要是输入到大模型调用的提示。然而，对于构建能够修改自身行为的自适应代理来说，只允许改变工作流的一部分而不允许其他部分可能会受到限制。在这个实验中，测试了Trace在联合提示优化和代码生成方面的性能。具体来说，优化了一个基于DSPy的大模型代理，并调整了其三个组件：元提示prompt_template、一个修改提示的函数create_prompt以及一个后处理大模型调用输出的函数extract_answer。实验设置了一个端到端的提示和代码优化管道，并使用自动评估函数将大模型的输出与真实情况相比较，这要求大模型代理不仅要生成正确答案，还要以正确的格式生成（遵循[30]的指导方针）。

实验使用了Big-Bench Hard数据集（15个示例用于训练，5个用于验证，其余用于测试）。实验比较了Trace与DSPy的COPRO模块（优化元提示）。表1显示，Trace能够优化DSPy程序，超越了DSPy的COPRO优化器的能力，特别是在算法任务上。这一结果展示了Trace如何具体改进现有的大模型提示优化库。

实验的第四部分测试了Trace优化长期工作流的能力，这些工作流具有复杂的依赖关系，并能够“通过时间进行反向传播”。

实验使用Trace训练模拟Sawyer机器人操纵器的控制器代码。实验使用了LLF-Bench的Meta-World环境作为模拟器，并考虑了三个任务：Reach、Pick-place和Push。对于每个任务，LLF-Bench提供了任务指令和动作空间的含义，这些被用作OPTO问题的上下文ω。观察是一个指示终端执行器位置、目标位置、夹爪状态等的Python字典向量。

动作空间是一个4维向量，用于控制终端执行器的相对位置和夹爪状态。在每个时间步骤中，LLF-Bench Meta-World模拟器返回观察结果和自然语言反馈以指导机器人。一个episode在机器人成功解决问题或超时后结束。

实验考虑了一个episodic训练设置。所有训练迭代的初始条件都是相同的。实验从10个保留的初始条件开始评估学习到的策略的成功与否。任务视野是10步，这对于任务完成是足够的，每个训练迭代有一个rollout。OPTO中的输出反馈是任务成功的文本表示。除了控制器代码，实验还装饰了gym环境的reset和step函数，以便整个rollout可以端到端地被追踪。

实验比较了Trace与OPRO；由于流式OPTO设置，OPRO的实现在每个迭代中只提出一个候选方案，然后对其进行评估并提供输出反馈。

图6总结了实验结果。Trace表示使用记忆的OptoPrime，而Trace NoMem表示不使用记忆的OptoPrime。OptoPrime显然是表现最好的优化器，特别是使用记忆的版本。OPRO能够在开始时解决Reach任务，但其性能随着迭代的进行而下降（这种不稳定性在[13]中提到），并且在Push任务中与不使用记忆的OptoPrime表现相似。为了验证OptoPrime的性能确实是由于使用了执行迹线，实验包括了一个消融实验，其中遮蔽了执行迹线，导致性能和稳定性显著下降。这次实验在所有实验中具有最复杂的图结构。实验结果非常令人印象深刻，表明Trace能够在数十次交互中学习复杂的控制逻辑，不仅适用于训练初始条件，还适用于保留的测试条件。

提交给NeurIPS后，另一个类似AutoDiff的框架TextGrad被发布，它与Trace有着相同的目标，即端到端优化AI工作流。相比之下，TextGrad传播文本反馈，而Trace传播最小子图。Trace的基于图的设计，将追踪基础设施和优化算法分开，使其更加灵活。实际上，实验中很容易将TextGrad实现为Trace框架中的一个优化器，但反过来则不可能（因为TextGrad将基础设施和优化算法耦合在一起）。

此外，与TextGrad不同，Trace支持联合优化异构参数，并且可以直接应用于给定的计算工作流的追踪，无需使用预定义的模板重写工作流。实验中将Trace应用于直接装饰TextGrad库发布的评估代码，并按照他们的训练/评估管道逐行优化参数。

这种实验设计通过允许每个优化器在相同时间访问相同的LLM API，使比较变得公平，并展示了Trace框架优化任何计算工作流的灵活性。实验选择了解决方案优化和提示优化作为推理任务实验。所有这些算法在这些实验中都取得了类似的成功率。一个值得注意的区别是，OptoPrime比TextGrad快约3倍的墙上时间，因为OptoPrime在每个优化步骤中只对LLM进行一次调用，而TextGrad的调用次数与图的大小成线性关系。

表2显示了Trace和TextGrad之间的比较结果。优化器是GPT-4o-2024-08-06，学生模型是GPT-35-turbo-1106。结果显示了最后迭代的成功率的平均值和标准误差，这些是由5个种子计算得出的。报告的实验时间以分钟为单位（时间不仅包括训练，还包括通过运行TextGrad的原始管道进行验证和测试）；GSM8K实验的时间被省略了，因为实验时间（>8小时）主要由评估而不是优化决定。

论文链接：https://arxiv.org/abs/2406.16218

项目链接：https://microsoft.github.io/Trace