[Paper Review] Joint Prompt Optimization of Stacked LLMs using Variational Inference

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

Topic Modeling5 citations

TL;DR

This paper proposes a variational inference framework for joint prompt optimization in stacked large language models (LLMs), treating hidden layer outputs as latent variables. By training prompts end-to-end using variational bounds, the method achieves performance comparable to GPT-4 using smaller, less powerful LLMs, demonstrating state-of-the-art results on reasoning and NLU tasks through a learnable, modular Deep Language Network (DLN) architecture.

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.

Motivation & Objective

• To address the high cost and limited adaptability of training large language models from scratch by enabling modular, reusable LLM components.
• To improve few-shot and zero-shot prompting by jointly optimizing prompts across stacked LLM layers.
• To formalize chain-of-thought and self-consistency prompting as structured probabilistic models using latent variable inference.
• To demonstrate that smaller LLMs can achieve high performance when stacked and prompted using a learnable, differentiable framework.
• To enable end-to-end training of multi-layer LLM pipelines through variational inference over hidden sequence outputs.

Proposed method

• Treats each LLM layer as a stochastic language layer where the prompt is a learnable parameter, forming a Deep Language Network (DLN).
• Models the output of the first LLM layer as a latent variable $ h $, enabling marginalization over $ h $ in a two-layer DLN (DLN-2).
• Applies variational inference to optimize the joint log-likelihood of the final output by approximating the true posterior over $ h $.
• Uses a variational posterior $ q(h|x, \pi_1) $ to approximate the true posterior $ p(h|x, \pi_0) $, enabling gradient-based optimization of both prompts $ \pi_0 $ and $ \pi_1 $.
• Employs a template-based input construction that combines input $ x $, prompt $ \pi $, and hidden states into a single sequence for LLM processing.
• Extends Automatic Prompt Engineer (APE) to include in-context examples and verbalized hard examples in the prompt, improving performance over baseline 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

Experimental results

Research questions

• RQ1 Can joint prompt optimization across stacked LLMs significantly improve performance on reasoning and NLU tasks?
• RQ2 Can a two-layer Deep Language Network (DLN-2) achieve performance comparable to GPT-4 using smaller, less powerful LLMs?
• RQ3 How can chain-of-thought and self-consistency prompting be formalized as a probabilistic graphical model with learnable priors?
• RQ4 Can variational inference over hidden LLM outputs enable effective, end-to-end training of multi-layer LLM pipelines?
• RQ5 Does learning task-specific prompt templates improve generalization and performance over fixed or hand-crafted prompts?

Key findings

• The proposed joint prompt optimization method outperforms standard APE on multiple reasoning and natural language understanding tasks by incorporating verbalized hard examples into prompts.
• DLN-2 achieves higher performance than single-layer models, demonstrating the benefit of stacking LLMs with optimized prompts.
• The method enables performance comparable to GPT-4 when using smaller, less powerful LLMs, suggesting a scalable alternative to training massive models.
• Chain-of-Thought (CoT) and Self-Consistency (SC-CoT) prompting are formalized as special cases of the DLN-2 framework with learnable priors over reasoning paths.
• The variational inference framework successfully trains both prompts $ \pi_0 $ and $ \pi_1 $ end-to-end, with the posterior approximation enabling effective gradient flow.
• The open-source DLN codebase enables reuse and extension of the framework for various downstream tasks and model architectures.