GReaTer: Gradients over Reasoning Makes Smaller Language Models Strong Prompt Optimizers

TL;DR

Prompt design is a bottleneck for LLM performance, especially for smaller models. GReaTer introduces a gradient-guided approach that uses loss gradients over reasoning to select prompt tokens for open-source LMs, avoiding reliance on proprietary LLMs. It delivers consistent gains across BBH, GSM8k, and FOLIO and shows strong transferability across model sizes and families, sometimes matching prompts tuned by GPT-4. This work demonstrates that gradients over reasoning can be a powerful signal for robust, scalable prompt optimization.

Abstract

The effectiveness of large language models (LLMs) is closely tied to the design of prompts, making prompt optimization essential for enhancing their performance across a wide range of tasks. Many existing approaches to automating prompt engineering rely exclusively on textual feedback, refining prompts based solely on inference errors identified by large, computationally expensive LLMs. Unfortunately, smaller models struggle to generate high-quality feedback, resulting in complete dependence on large LLM judgment. Moreover, these methods fail to leverage more direct and finer-grained information, such as gradients, due to operating purely in text space. To this end, we introduce GReaTer, a novel prompt optimization technique that directly incorporates gradient information over task-specific reasoning. By utilizing task loss gradients, GReaTer enables self-optimization of prompts for open-source, lightweight language models without the need for costly closed-source LLMs. This allows high-performance prompt optimization without dependence on massive LLMs, closing the gap between smaller models and the sophisticated reasoning often needed for prompt refinement. Extensive evaluations across diverse reasoning tasks including BBH, GSM8k, and FOLIO demonstrate that GReaTer consistently outperforms previous state-of-the-art prompt optimization methods, even those reliant on powerful LLMs. Additionally, GReaTer-optimized prompts frequently exhibit better transferability and, in some cases, boost task performance to levels comparable to or surpassing those achieved by larger language models, highlighting the effectiveness of prompt optimization guided by gradients over reasoning. Code of GReaTer is available at https://github.com/psunlpgroup/GreaTer.

Figures (8)

• Figure 1: Comparison of textual feedback-based prompt optimization and GReaTer. Left: textual feedback relies entirely on a larger language model's judgments. Right: GReaTer avoids external large, proprietary models, using token suggestions from a small model and guiding prompt token selection with loss gradients. GReaTer incorporates model reasoning by first generating reasoning, then applying an extraction prompt to obtain answer logits for computing loss gradients. This "gradient over reasoning" approach optimizes using direct signals rather than relying on language model feedback.
• Figure 2: Overall workflow of GReaTer. (i) The language model $f_{\text{LLM}}$ generates token candidates by conditioning on input samples. (ii) $f_{\text{LLM}}$ uses task input and current prompt to generate reasoning and extract final answer logits. (iii) The logits are used to calculate loss and compute gradient over generated reasoning with respect to the candidate tokens. These gradients determine the selection of candidate token to update the current position of the current prompt.
• Figure 3: Ablation study on "Gradient Over Reasoning" in GReaTer. Gradient calculation without reasoning causes notable performance drops, showing the importance of reasoning for gradients.
• Figure 4: Efficacy of GReaTer in zero-shot setting compared to five-shot inference with Llama-3-8B-Instruct.
• Figure 5: Win/Draw/Loss Comparison of GReaTer and SOTA prompt optimization techniques APO, TextGrad, APE, and PE2 in optimization with Llama-3-8B-Instruct. GReaTer maintains a significant winning margin over these methods, highlighting its effectiveness in optimization.