Scaling Textual Gradients via Sampling-Based Momentum
2025/05/31 by Ding, Zixin, Hong, Junyuan, Shi, Zhan +6
#Artificial Intelligence (cs.AI) #Computation and Language (cs.CL) #FOS: Computer and information sciences
paper · doi:10.48550/arxiv.2506.00400
Abstract
LLM-based prompt optimization, that uses LLM-provided "textual gradients" (feedback) to refine prompts, has emerged an effective method for automatic prompt engineering. However, its scalability and stability are unclear when using more data in training. We systematically investigate the potential and challenges of scaling training data in textual gradient descent. We show that naively scaling training examples is infeasible due to both explicit context-length limits and an implicit context wall, where long-context degradation yields diminishing returns. Inspired by prior wisdom in stochastic gradient descent, we propose Textual Stochastic Gradient Descent with Momentum (TSGD-M), which reweights updates through momentum sampling, using bootstrapped minibatch validation accuracy as importance weights over historical prompts. We introduce Gumbel-Top-k sampling for prompt generation, balancing exploration--exploitation and improving sampling efficiency while maintaining a low-variance running mean estimator. TSGD-M integrates seamlessly into existing prompt optimization frameworks, including TextGrad, DSPy-COPRO, and AdalFlow, and achieves consistent gains across 5 benchmarks.
Citations
- Attention Sinks and Compression Valleys in LLMs are Two Sides of the Same Coin
- Context Length Alone Hurts LLM Performance Despite Perfect Retrieval
- LAD-VF: LLM-Automatic Differentiation Enables Fine-Tuning-Free Robot Planning from Formal Methods Feedback
- GEPA: Reflective Prompt Evolution Can Outperform Reinforcement Learning
- Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities
- S*: Test Time Scaling for Code Generation
- Multi-Agent Design: Optimizing Agents with Better Prompts and Topologies
- s1: Simple test-time scaling
- LLM-AutoDiff: Auto-Differentiate Any LLM Workflow
- Squeezed Attention: Accelerating Long Context Length LLM Inference
- Introducing MAPO: Momentum-Aided Gradient Descent Prompt Optimization
- How Susceptible are LLMs to Influence in Prompts?
- Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters
- The Llama 3 Herd of Models
- Large Language Monkeys: Scaling Inference Compute with Repeated Sampling
- Private prediction for large-scale synthetic text generation
- Enhancing In-Context Learning via Implicit Demonstration Augmentation
- Teach Better or Show Smarter? On Instructions and Exemplars in Automatic Prompt Optimization
- Can LLMs Learn by Teaching for Better Reasoning? A Preliminary Study
- Can Long-Context Language Models Subsume Retrieval, RAG, SQL, and More?
- Optimizing Instructions and Demonstrations for Multi-Stage Language Model Programs
- Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference
- BABILong: Testing the Limits of LLMs with Long Context Reasoning-in-a-Haystack
- TextGrad: Automatic "Differentiation" via Text
- Uncertainty Estimation and Quantification for LLMs: A Simple Supervised Approach
- Many-Shot In-Context Learning
- Differentially Private Synthetic Data via Foundation Model APIs 2: Text
- Same Task, More Tokens: the Impact of Input Length on the Reasoning Performance of Large Language Models
- InternLM-Math: Open Math Large Language Models Toward Verifiable Reasoning
- SGLang: Efficient Execution of Structured Language Model Programs
- When does In-context Learning Fall Short and Why? A Study on Specification-Heavy Tasks
- PromptAgent: Strategic Planning with Language Models Enables Expert-level Prompt Optimization
- Quantifying Language Models' Sensitivity to Spurious Features in Prompt Design or: How I learned to start worrying about prompt formatting
- Efficient Streaming Language Models with Attention Sinks
- Privacy-Preserving In-Context Learning with Differentially Private Few-Shot Generation
- Large Language Models as Optimizers
- Lost in the Middle: How Language Models Use Long Contexts
- Joint Prompt Optimization of Stacked LLMs using Variational Inference
- Automatic Prompt Optimization with "Gradient Descent" and Beam Search
- Cost-Effective Hyperparameter Optimization for Large Language Model Generation Inference
- Structured Prompting: Scaling In-Context Learning to 1,000 Examples
- Large Language Models Are Human-Level Prompt Engineers
- Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them
- Large Language Models are Zero-Shot Reasoners
- Training Compute-Optimal Large Language Models
- Training language models to follow instructions with human feedback
- Training Verifiers to Solve Math Word Problems
- Stochastic Batch Acquisition: A Simple Baseline for Deep Active Learning
- Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity
- Measuring Mathematical Problem Solving With the MATH Dataset
- An Improved Analysis of Stochastic Gradient Descent with Momentum
- Bandit Algorithms
- Language Models are Few-Shot Learners
- Scaling Laws for Neural Language Models
- Learning to Few-Shot Learn Across Diverse Natural Language\n Classification Tasks
- Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement
- HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering
- A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay
- Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge
- Attention Is All You Need
- A* Sampling
- A Sentimental Education: Sentiment Analysis Using Subjectivity Summarization Based on Minimum Cuts
Related