GFlowPO: Generative Flow Network as a Language Model Prompt Optimizer

TL;DR

GFlowPO presents a novel framework for automatic prompt optimization by casting prompt search as posterior inference over latent prompts, regularized by a meta‑prompt prior from a reference LM. It combines off‑policy GFlowNet training to amortize posterior sampling with a training‑free Dynamic Memory Update that updates the meta‑prompt based on diverse and high‑reward prompts, progressively concentrating search on high‑reward regions. The method demonstrates strong, robust gains across few‑shot text classification, instruction induction, and QA benchmarks, outperforming on‑policy baselines like StablePrompt and other discrete prompt strategies. The approach highlights the value of memory‑augmented, off‑policy exploration in discrete prompt spaces and opens avenues for more sample‑efficient, scalable prompt optimization in LMs.

Abstract

Finding effective prompts for language models (LMs) is critical yet notoriously difficult: the prompt space is combinatorially large, rewards are sparse due to expensive target-LM evaluation. Yet, existing RL-based prompt optimizers often rely on on-policy updates and a meta-prompt sampled from a fixed distribution, leading to poor sample efficiency. We propose GFlowPO, a probabilistic prompt optimization framework that casts prompt search as a posterior inference problem over latent prompts regularized by a meta-prompted reference-LM prior. In the first step, we fine-tune a lightweight prompt-LM with an off-policy Generative Flow Network (GFlowNet) objective, using a replay-based training policy that reuses past prompt evaluations to enable sample-efficient exploration. In the second step, we introduce Dynamic Memory Update (DMU), a training-free mechanism that updates the meta-prompt by injecting both (i) diverse prompts from a replay buffer and (ii) top-performing prompts from a small priority queue, thereby progressively concentrating the search process on high-reward regions. Across few-shot text classification, instruction induction benchmarks, and question answering tasks, GFlowPO consistently outperforms recent discrete prompt optimization baselines.

Figures (7)

• Figure 1: Concepts. Blue contour indicates high performing prompt regions. (a) Existing on-policy RL frameworks fail to explore the huge combinatorial search space with poor sample efficiency. (b) Our GFlowPO that can sample efficiently explore the search space by gradually annealing the posterior density (dotted ellipses) towards more promising area with the off-policy GFlowNet training in step-A and DMU in step-B.
• Figure 2: GFlowPO pipeline. The optimizer prompt-LM samples prompts conditioned on meta-prompt $M$, the target LLM provides rewards, and off-policy GFlowNet training plus Dynamic Memory Update (DMU) iteratively improves exploration and prompt quality.
• Figure 3: Meta-prompt template used in GFlowPO.
• Figure 4: Heatmap of performances on few-shot text classification tasks (6 datasets) with various prompt-LM, target LM pairs. Each number is averaged over six tasks. MP: Manual Prompt, G2: Gemma-2B, G7: Gemma-7B, M7: Mistral-7B, L8: Llama-3-8B, F11: Falcon-11B. GFlowPO works well across various (prompt-LM, target LM) pairs.
• Figure 5: Training accuracy as a function of the number of sampled prompts on four II tasks. Solid lines denote the mean training accuracy over three random seeds, and shaded regions indicate one standard deviation. GFlowPO consistently discovers higher-reward prompts more efficiently than StablePrompt as sampling progresses.