[论文解读] Joint Prompt Optimization of Stacked LLMs using Variational Inference

Alessandro Sordoni, Xingdi Yuan+7 more|arXiv (Cornell University)|Jun 21, 2023

Topic Modeling被引用 5

一句话总结

本文提出了一种变分推断框架，用于在堆叠的大规模语言模型（LLMs）中进行联合提示优化，将隐藏层输出视为潜在变量。通过使用变分界端到端训练提示，该方法在较小、计算能力较弱的 LLM 上实现了与 GPT-4 相当的性能，在推理和自然语言理解（NLU）任务上展现出最先进结果，其核心是可学习的、模块化的深度语言网络（DLN）架构。

ABSTRACT

Large language models (LLMs) can be seen as atomic units of computation mapping sequences to a distribution over sequences. Thus, they can be seen as stochastic language layers in a language network, where the learnable parameters are the natural language prompts at each layer. By stacking two such layers and feeding the output of one layer to the next, we obtain a Deep Language Network (DLN). We first show how to effectively perform prompt optimization for a 1-Layer language network (DLN-1). Then, we present an extension that applies to 2-layer DLNs (DLN-2), where two prompts must be learned. The key idea is to consider the output of the first layer as a latent variable, which requires inference, and prompts to be learned as the parameters of the generative distribution. We first test the effectiveness of DLN-1 in multiple reasoning and natural language understanding tasks. Then, we show that DLN-2 can reach higher performance than a single layer, showing promise that we might reach comparable performance to GPT-4, even when each LLM in the network is smaller and less powerful.

研究动机与目标

• 为解决从零开始训练大规模语言模型的高昂成本和有限适应性，通过实现模块化、可重用的 LLM 组件来达成目标。
• 通过在堆叠的 LLM 层之间联合优化提示，提升少样本和零样本提示性能。
• 将思维链（chain-of-thought）和自一致性提示形式化为使用潜在变量推断的结构化概率模型。
• 证明当使用可学习的、可微分的框架对堆叠的 LLM 进行提示时，小型 LLM 也能实现高性能。
• 通过隐藏序列输出上的变分推断，实现多层 LLM 管道的端到端训练。

提出的方法

• 将每个 LLM 层视为一个随机语言层，其中提示是可学习参数，从而构成深度语言网络（DLN）。
• 将第一层 LLM 的输出建模为潜在变量 $ h $，在两层 DLN（DLN-2）中实现对 $ h $ 的边际化。
• 应用变分推断，通过近似 $ h $ 的真实后验分布来优化最终输出的联合对数似然。
• 使用变分后验 $ q(h|x, \pi_1) $ 近似真实后验 $ p(h|x, \pi_0) $，从而实现对提示 $ \pi_0 $ 和 $ \pi_1 $ 的基于梯度的优化。
• 采用基于模板的输入构建方式，将输入 $ x $、提示 $ \pi $ 和隐藏状态整合为单个序列以供 LLM 处理。
• 将自动提示工程师（APE）扩展为包含上下文示例和显式化困难样本的提示，从而在性能上超越基线 APE。

Figure 1: Left : An illustration of a DLN-1 performing a sentiment analysis task: input and the trainable prompt are merged using a template and fed to the LM for answer generation. Right : a DLN-2 with a residual connection, performing the date understanding task: two prompts need to be learned. In

实验结果

研究问题

• RQ1 在堆叠的 LLM 中进行联合提示优化，能否显著提升推理和 NLU 任务的性能？
• RQ2 两层深度语言网络（DLN-2）能否在使用较小、计算能力较弱的 LLM 时，实现与 GPT-4 相当的性能？
• RQ3 如何将思维链（CoT）和自一致性提示形式化为具有可学习先验的概率图模型？
• RQ4 对 LLM 隐藏输出进行变分推断，能否实现对多层 LLM 管道的有效、端到端训练？
• RQ5 学习特定任务的提示模板是否能提升泛化能力与性能，相比固定或人工设计的提示？

主要发现

• 所提出的联合提示优化方法通过在提示中引入显式化困难样本，在多个推理和自然语言理解任务上优于标准 APE。
• DLN-2 的性能高于单层模型，证明了通过优化提示堆叠 LLM 的有效性。
• 当使用较小、计算能力较弱的 LLM 时，该方法实现了与 GPT-4 相当的性能，表明其为训练超大模型提供了一种可扩展的替代方案。
• 思维链（CoT）和自一致性（SC-CoT）提示可形式化为 DLN-2 框架的特例，其推理路径具有可学习先验。
• 变分推断框架成功实现了对提示 $ \pi_0 $ 和 $ \pi_1 $ 的端到端训练，后验近似确保了有效的梯度传播。
• 开源的 DLN 代码库支持在各种下游任务和模型架构中复用与扩展该框架。