PaperlayerView exact revision on arXiv

Source-linked AI summary

Joint Prompt Optimization of Stacked LLMs using Variational Inference

Alessandro Sordoni, Xingdi Yuan, Marc-Alexandre Côté, Matheus Pereira, Adam Trischler, Ziang Xiao, Arian Hosseini, Friederike Niedtner, Nicolas Le Roux

arXiv:2306.12509v2cs.CLcs.LG

TL;DR

The paper addresses how to improve smaller LLMs and compose them without training their weights by treating prompts as parameters of stacked stochastic language layers. It develops local-search prompt optimization for DLN-1 and variational inference for DLN-2, where the first layer’s output is latent. The reported results show gains from stacking and prompting smaller LLMs, while the framework remains limited by engineered templates and evaluation of only one- and two-layer networks.

  • Problem

    The paper studies how to obtain useful computation from smaller LLMs and compose existing models while avoiding the technical and societal costs of continually increasing model size.

  • Method

    It models prompted LLMs as stochastic language layers, optimizes one-layer prompts by extending APE, and trains two-layer networks with variational inference over the first layer’s latent output.

  • Results

    The method improves DLN-1 downstream performance over APE on several tasks, and stacked prompted smaller LLMs show boosted performance compared with a single layer.

  • Takeaways & Limitations

    Prompt-based stacking provides a modular way to improve smaller LLMs and reuse multiple existing models across new use cases.

  • Takeaways & Limitations

    The study tests only one- and two-layer language networks, uses engineered backward and forward templates, and notes that template learning remains future work.

Abstract

from arXiv · show

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.

1 Introduction

The paper frames LLMs as stochastic language layers whose prompts are learnable parameters, then stacks them into Deep Language Networks. It extends prompt optimization from one layer to two layers by treating the first layer’s output as a latent variable and using variational inference.

  • DLN framework: LLMs are modeled as stochastic language layers, and stacking layers with prompt-controlled templates forms a Deep Language Network.Each layer maps input strings to output distributions, with prompts serving as learnable parameters.
  • DLN-1: DLN-1 prompt optimization extends APE by combining instruction directives with verbalized task examples, surpassing APE on several tasks.The approach targets reasoning and natural language understanding tasks.
  • DLN-2: DLN-2 treats the first layer’s output h as a latent variable and uses variational inference to optimize prompts for the two-layer generative distribution.The output layer remains conditioned on the original input through a residual connection.
  • Connections to prompting: Viewing hidden outputs as latent variables connects the framework to CoT and self-consistency, while learning a task-specific prior over successful CoT strings.CoT can be represented as a DLN-2 with fixed first- and second-layer prompts; the framework can learn supplements or replacements.
  • Scope: The paper positions stacked prompted LLMs as a route toward improving smaller models while retaining modularity across reasoning and language-understanding tasks.The stated program is to train DLN-2 systems and evaluate them across these task families.

2 One-Layer Language Networks

The one-layer network treats a prompt as the parameter controlling an LLM’s output distribution and optimizes it through discrete local search. Candidate prompts are proposed by an LLM, scored by output likelihood, and updated iteratively.

  • Language-layer formulation: A language layer defines pLM(y|F(x, π)), where a template combines input x and prompt π before the LLM generates output y.Templates specify how prompts and inputs are organized into one sequence.
  • Optimization strategy: Prompt optimization uses discrete local search because the prompt space is not suitable for gradient-based optimization.The method extends Automatic Prompt Engineer with LLM-generated local prompt modifications.
  • Prompt proposal: Each iteration samples N candidate prompts conditioned on a minibatch, predictions, targets, and the current prompt, using a backward proposal template.The proposal distribution is implemented by an LLM and is designed to generate diverse candidate modifications.
  • Prompt selection: Candidates are scored by normalized output log-likelihood, and the highest-scoring prompt is selected with backtracking and a memory of strong prompts improving robustness.The selection rule is π = arg maxπn log pLM(y|F(x; πn)), normalized by output length in practice.
  • Outcome: The resulting DLN-1 prompts combine zero-shot-style directives with in-context task examples and improve downstream performance over APE on several tasks.The one-layer results are reported in Table 1 before extending the method to two layers.

3 Two-Layer Deep Language Networks (DLN-2)

DLN-2 stacks two language layers, treating the first layer’s output as a latent variable and learning both prompts through variational inference. This formulation increases expressiveness but makes joint prompt optimization harder.

  • 3 Two-Layer Deep Language Networks (DLN-2): DLN-2 feeds the first language layer’s output into the second, with a residual connection from the original input to the output layer.The intermediate output h is treated as a latent variable that may make the target y easier to explain.
  • 3 Two-Layer Deep Language Networks (DLN-2): Jointly optimizing the two prompts is impractical with random search and exponentially harder to tune manually than a single prompt.The paper therefore turns to variational inference.
  • 3.1 Variational Inference Objective: Variational inference introduces an approximate posterior q(h) and an ELBO that decomposes optimization into separate problems for π0 and π1.The π1 search uses posterior samples as additional inputs, while π0 is optimized using those samples as targets.
  • 3.1 Variational Inference Objective: The approximation is useful only when q(h) is close to the true posterior, so the method applies posterior sharpening to reduce the KL gap.Posterior sharpening reweights hidden-state samples using their likelihood under the forward model.
  • 3.1 Variational Inference Objective: The hidden proposal uses an LLM to sample candidate states from prior, label-informed, or edited distributions, with a mixture of qpri and qpri+ found most effective.The edited proposal rewrites a forward-pass hidden state using the target label and second-layer prompt.

4 Practical Instantiation

The practical DLN-2 implementation uses LLM-generated proposals, posterior-weighted hidden states, and alternating local search over both prompts. Additional engineering strategies address proposal diversity, memory, and slow first-layer updates.

  • 4 Practical Instantiation: Proposal diversity is increased by varying backward-template instructions and asking the LLM to provide useful examples for prompt improvement.These strategies use pools of hand-written meta-instructions.
  • 4 Practical Instantiation: The training algorithm initializes both prompts, samples hidden-state proposals, computes prior and output likelihoods, and selects the highest-weight hidden state.Posterior weights are normalized from the sum of prior log-probability and output log-likelihood.
  • 4 Practical Instantiation: Candidate prompts for π0 and π1 are generated with backward templates and selected by their scores during alternating optimization.The proposal process uses task batches, outputs, hidden states, and the current prompt as conditioning information.
  • 4 Practical Instantiation: Prompt proposals can contain synthetic task examples, which often outperform standard in-context learning in the authors’ observations.The authors hypothesize that verbalizing examples and selecting them according to training errors improves their usefulness.
  • 4 Practical Instantiation: Backtracking retains the current prompt among candidates so sampling or selection errors need not force an update.This is necessary because gradient information is unavailable and candidate sets are restricted for computational reasons.
  • 4 Practical Instantiation: An exploration reward is added because shared posterior and prior templates can cause the first-layer prompt π0 to update too slowly.The reward favors candidates associated with hidden states that previously produced incorrect predictions.

5 Experiments and Results

Experiments evaluate whether DLN-1 improves over prompt baselines and whether depth improves over DLN-1 across nine classification and reasoning tasks. DLN-2 provides its largest gains on several tasks where DLN-1 is weaker.

  • 5 Experiments and Results: The experiments address whether DLN-1 can outperform APE and in-context learning and whether additional network depth improves performance.The study uses nine NLP and reasoning classification tasks, accuracy, GPT-3 and GPT-4 backbones, and several baseline settings.
  • 5.2 DLN-1: On Hyperbaton, DLN-1 prompts combine instructions with automatically selected training examples, linking prompt optimization with in-context learning.The selected examples are chosen for their impact on performance and resemble a combination of KATE-style example selection and APE-style prompt selection.
  • 5.2 DLN-1: DLN-1 significantly beats the best GPT-3 baseline by about 20, 10, and 7 percentage points on Hyperbaton, TREC, and Subjectivity, respectively.It also surpasses GPT-4 baselines on Hyperbaton, TREC, and Disaster but underperforms GPT-4 on the other tasks.
  • 5.2 DLN-1: DLN-1 matches the best GPT-3-based method on Disaster, MPQA, and Airline and narrowly beats the best GPT-3 baseline on Logic.7 and Navigation.Table 1 reports test accuracy averaged over three random seeds, with confidence intervals for trainable systems and GPT-4.
  • 5.2 DLN-1: DLN-1 can underperform the zero-shot baseline on Date because the authors observed overfitting caused by too few validation examples.This limitation is reported specifically for the Date task.
  • 5.3 DLN-2: DLN-2 improves over DLN-1 by 7.2% absolute on average, with large gains on Navigation and Date and further improvements on Subjectivity and Disaster.On Navigation and Date, DLN-2 outperforms all single-layer networks; on Logic.7, methods perform similarly.

6 Related Work

Prior work studies prompting, prompt optimization, and compositions of LLMs, while DLNs extend these ideas with learned multi-layer inference applied to downstream tasks.

  • Prompt-Based Machine Learning: Chain-of-thought prompting improves LLM performance by eliciting intermediate reasoning steps, either through examples or an explicit step-by-step instruction.
  • Prompt Optimization Techniques: Automatic prompt optimization includes heuristic methods such as Autoprompt and GRIPS, while later work uses reinforcement learning to explore discrete prompt spaces more systematically.
  • Multi-Layer LLM systems: Composed-LLM systems include control-flow generation, interactive AI chains, and language model cascades that represent LLM compositions as graphical models with string-valued variables.
  • Multi-Layer LLM systems: DLNs go beyond the conceptual cascade framework by providing an inference technique for LLM-based graphical models and applying learned LLM networks to downstream tasks.

7 Conclusion and Future Work

The paper introduces joint prompt optimization for stacked LLMs by treating hidden-layer outputs as latent variables, and reports improved performance and modularity while identifying deployment and scaling limits.

  • Conclusion: Joint prompt optimization learns prompts in deep LLM networks by performing inference over outputs of hidden language layers.
  • Conclusion: CoT, Generated Knowledge Prompting, ReAct, Reflexicon, and Self-Consistency can be interpreted as special cases or ensembles within the DLN framework.
  • Conclusion: Stacking and properly prompting smaller LLMs boosts performance, while modularity may improve adaptation and reuse across use cases.
  • Future Work: The study tests only one- and two-layer networks, relies on engineered templates, and leaves template learning and broader graph structures for future work.
  • Impact and limitations: The evaluation uses artificial benchmarks and does not establish when deployment is appropriate, performance in uncontrolled environments, or safeguards against misuse.

Contents in Appendices:

The appendices document authorship, experimental details, prompt initializations, additional experiments, templates, and extensions to multilayer training.

  • Appendix B: Appendix B provides task statistics and the prompt strings used to initialize DLNs.
  • Appendices C–D: Appendix C reports additional experiments and baseline comparisons, while Appendix D presents the forward and backward templates used in DLN.
  • Appendix E: Appendix E gives an algorithm that generalizes DLN training to multiple layers.
  • Author Contributions: Author contributions cover the DLN idea, variational inference formulation, prompt optimization, experiments, infrastructure, templates, evaluation, and project management.Contributors also supported code, demos, literature review, and paper writing.

B.1 Additional Task Information

This appendix documents the benchmark tasks and the prompt initializations used for the one- and two-layer DLN experiments.

  • Additional Task Information: The study uses nine NLP and reasoning classification tasks drawn from BigBench-Hard, BigBench, and prior zero- or few-shot learning benchmarks.
  • Additional Task Information: Table 3 lists the tasks used in the experiments and their statistics.
  • Prompt Initialization: Table 4 records prompt initializations, including task descriptions for classification layers and selected hidden-layer prompts for two-layer DLNs.The hidden layer uses “Decompose the problem to make it simpler:” for Nav. and Subj., and an empty string for Date. and Logic.

C Additional Experiments

Additional experiments compare DLN variants with prompting baselines, open-source models, and alternative DLN-2 training strategies. They also document templates, hidden-state proposals, and generalized multi-layer training procedures.

  • Comparisons: DLN-1 and DLN-2 outperform CoT+APE on GPT-3 across the reported additional experiments.Table 5 reports test accuracy averaged over three seeds with 95% confidence intervals.
  • Comparisons: Increasing ICL examples improves performance but generally does not match DLN-1 and DLN-2, while context length limits ICL and KATE 32-shot.These comparisons are reported in Table 6 using GPT-3.
  • Open-source models: DLN-1 outperforms ICL on all reported tasks when using WizardLM-v1.2 13B, whose few-shot context learning is described as weaker.The open-source-model comparison appears in Table 7.
  • DLN-2 training: The DLN-2 layer-wise alternative initializes the first-layer prompt from DLN-1, then learns the bottom-layer prompt with variational inference, with or without fine-tuning the last layer.Table 9 compares layer-wise and end-to-end DLN-2 training.
  • Templates and training: The method proposes prompts from task examples and model predictions, samples hidden states through forward and answer-conditioned backward templates, and can generalize training to multiple layers.The backward proposal can generate synthetic examples in prompts, while the generalized algorithm iterates over layers and posterior samples.

I Implementation Details

Implementation details specify the hyperparameter search, prompt-ranking controls, and evaluation protocol used for the experiments.

  • Hyperparameter search: The experiments search hyperparameters including backward-template type, exploration penalty, hidden-state sample count, memory, held-out ranking, and validation tolerance.The search space is summarized in Table 10.
  • Hyperparameter search: The v3.5 backward template generally works better because its “Be concise.” instruction regularizes the length of discovered prompts.The authors note that more principled length regularization remains future work.
  • DLN-2 settings: Two-layer experiments restrict the search space because of computational costs, using v3.5, tolerance 2, memory 2, held-out ranking, and logp_penalty 0.5.These settings are reported specifically for DLN-2.
  • Evaluation: Evaluation uses exact-match accuracy averaged over three random seeds after output normalization for tokenization and capitalization.The reported test accuracy and inference costs are summarized in Table 11.

J Pricing

The pricing analysis quantifies both API usage and deployment-oriented test-time inference cost for DLN experiments.

  • Training cost: A GPT-3 run costs roughly $59 for DLN-1 and $273 for DLN-2, based on 2,941,360 and 13,654,962 tokens respectively.The calculation uses a GPT-3 price of $0.02 per 1,000 tokens and Hyperbaton as an example.
  • Test-time cost: DLN-1 improves over ICL on 5 of 9 GPT-4 tasks at a comparable test-time token cost.Table 11 emphasizes testing cost because training cost is one-off for real-world deployment.
  • Test-time cost: Test-time costs vary by task: reasoning tasks may not benefit from ICL, whereas Subj, Trec, and Hyper. benefit significantly.The comparison is reported in the Table 11 pricing discussion.
Loading 2306.12509v2…