Composing Policy Gradients and Prompt Optimization for Language Model Programs

1 Introduction

• Multi-module programs separate subtasks across LM calls, routing structured outputs between modules while keeping raw generation traces local.This modularity supports controllability, testing, parallelization, and prompt-level optimization.Source: §1 · p. 1, Figure 1 · p. 2
• GRPO was designed for single-stage rollouts, leaving open how to apply it when modular trajectories differ in context, length, control flow, or intermediate state.Different rollouts can invoke different modules, terminate early, or fail during parsing.Source: §1 · pp. 1–2
• mmGRPO adapts GRPO to arbitrary multi-module programs by forming groups from structurally comparable module invocations, with trajectory-level grouping as an alternative.Both variants extend existing prompt-optimization abstractions and are designed as interchangeable off-the-shelf optimizers.Source: §1 · p. 2
• The evaluation compares prompt optimization, mmGRPO, and their combination across classification, multi-hop claim verification, and privacy-conscious delegation using two open-source LMs.The tasks come from a subset of LangProBe and use llama3.1-8b-instruct and qwen3-8b.Source: §1 · p. 2
• 11%: BetterTogether combining MIPROv2 and mmGRPO improves average performance over the model’s unadapted reasoning performance, while also gaining 5% over MIPROv2 alone.The combined approach also improves by 3% over mmGRPO alone.Source: §1 · p. 2

2 Preliminaries

• GRPO groups same-prompt trajectories, normalizes their scalar rewards into advantages, and upweights relatively higher-reward completions with PPO-style stabilization.The normalized advantage is applied uniformly across all tokens in a completion.Source: §2 · pp. 2–3
• An LM program orchestrates module invocations and tools, transforming structured inputs into final outputs and trajectories of module calls.Each trajectory trace records a module’s identity, inputs, and outputs in execution order while omitting other control logic.Source: §2 · p. 3
• Each module has a prompt template and LM weights, may be invoked repeatedly, and returns sampled outputs to the program’s control flow.Repeated query-generation calls can occur across multiple hops within one trajectory.Source: §2 · p. 3
• Modular control flow avoids the context growth and information leakage of single-prompt ReAct-style agents while making subtasks independently testable.This explicit modularity is why single-expanding-prompt GRPO implementations do not directly apply.Source: §2 · p. 3
• LM-program optimization maximizes reward from executions, with prompt optimization updating templates, mmGRPO updating module weights, and BetterTogether composing both updates.The reward can use metadata such as gold answers for evaluation even though that metadata is hidden during execution.Source: §2 · p. 3

Framework design

• mmGRPO remains interchangeable with prompt optimizers because it uses existing program execution traces without modifying program logic.Learning algorithms can therefore be swapped without entangling program design with optimizer-specific details.Source: §3 · p. 3

Training algorithm

• The shared training loop samples batches and rollouts, forms groups paired with module weights, and applies a GRPO-style update to each group.Both variants share this control flow and differ only in their grouping function and objective.Source: §3 · pp. 4–5

Module-level grouping

• Module-level mmGRPO forms groups by aligning calls with the same module identifier and relative invocation index across sampled trajectories.Each aligned entry retains its local prompt, module output, and trajectory-level reward.Source: §3 · p. 4
• Divergent trajectories require padding when different module calls, early termination, or parsing failures produce groups smaller than the target size.The padding strategy keeps group sizes fixed before applying the loss.Source: §3 · p. 4, Figure 2 · p. 4

Module-level objective

• Module-level groups may contain different local prompts and can update separate module weights or a shared adapter through independently applied group losses.The experiments use one shared LM adapter for all trainable modules.Source: §3 · p. 5

Grouping variants

• Module-level and trajectory-level mmGRPO differ in normalization: module-level uses each invocation’s output length, whereas trajectory-level uses the full rollout’s generated length.Trajectory-level grouping avoids module-specific alignment and padding assumptions but keeps the entire execution as the grouped unit.Source: §3 · p. 5

Scope and extensions

• The framework currently samples rollouts on-policy from the student program, while teacher-program mixtures for partially off-policy training remain future work.Teacher programs could differ in prompt templates or LM weights while sharing the student’s structural interface.Source: §3 · p. 5

BetterTogether

• BetterTogether first optimizes prompt templates and then applies weight optimization to the resulting program.The composition targets different program elements: prompts shape rollouts, while weight optimization updates the model under those induced rollouts.Source: §4 · p. 6

Online RL composition

• This work extends BetterTogether to online RL by applying module-level mmGRPO after automatic prompt optimization with MIPROv2.The authors focus on this ordering rather than prompt optimization after weight tuning.Source: §4 · p. 6

5.1 LMs and datasets

• The evaluation covers Banking77 intent classification, PAPILLON privacy-conscious delegation, and HoVer4-HOP multi-hop retrieval using baseline, prompt-optimized, mmGRPO, and combined configurations.Banking77 uses a single CoT module as a sanity check, PAPILLON has redaction and response-synthesis modules, and HoVer4-HOP iterates query generation and fact summarization over four hops.Source: §5.1 · p. 6, §5.2 · p. 7

5.2 Baseline and method details

• MIPROv2 jointly optimizes module instructions and optional few-shot examples, while BetterTogether fixes those prompts before mmGRPO optimizes program weights from online full-execution rewards.Neither method uses intermediate module-level supervision.Source: §5.2 · p. 7

5.3 Main results

• 7%: mmGRPO improves over Vanilla CoT on average, while BetterTogether improves over MIPROv2 by 5% across LM–task pairs.The results support gains from both weight optimization and composing prompt optimization with mmGRPO.Source: §5.3 · p. 7, Table 1 · p. 8
• MIPROv2 improves Vanilla CoT by 6% versus mmGRPO’s 7% but requires 1.4 hours and one H100 GPU on average, compared with 18.7 hours and two H100 GPUs for mmGRPO.Prompt optimization is therefore more feasible under lower computation budgets in these experiments.Source: §5.3 · p. 7
• 11%: BetterTogether(PO, mmGRPO) improves over Vanilla CoT, with additional gains of 5% over MIPROv2 and 3% over vanilla mmGRPO.Table 1 reports evaluation accuracy averaged over three seeds across Banking77, PAPILLON, and HoVer4-HOP.Source: §5.3 · p. 8, Table 1 · p. 8

5.4 Trajectory-level mmGRPO experiment

• Trajectory-level mmGRPO reaches Recall@100 of 75.3 versus 71.0 for module-level mmGRPO on HoVer with CISPO and additional group and batch-size tuning.This variant groups full program executions rather than aligned module invocations.Source: §5.4 · p. 8

5.5 Qualitative case studies

• On HoVer, mmGRPO produces shorter, more targeted retrieval queries and sometimes passes forward whether evidence has actually been established.One query retrieves all three relevant documents versus two for MIPROv2 and zero for Vanilla CoT; another case records that the Manning Field claim remains unsupported.Source: §5.5 · p. 8

Prompt optimization

• Prompt optimization adapts LM prompts through instruction generation, gradient-based discrete-text methods, or reinforcement-learning editors.Source: §6 · p. 9

Weight optimization

• PPO supports reinforcement-learning post-training but requires both policy and reward models, making it computationally intensive.Source: §6 · p. 9
• DPO and GRPO simplify PPO by avoiding explicit reward modeling or value models, respectively, while using preference pairs or relative group rewards.Source: §6 · p. 9

Optimization of LM programs’ prompts and weights

• Existing LM-program optimization includes prompt optimizers and offline combinations of weight and prompt optimization, but online GRPO for multi-module programs remains an implementation challenge.Source: §6 · p. 9

7 Conclusion

• mmGRPO propagates trajectory rewards across disjoint modules and can bolt onto existing programs while accommodating heterogeneous prompts and partial trajectories.The framework is designed to avoid intermediate supervision and preserve separation between programming and learning concerns.Source: §7 · p. 9
• Combining mmGRPO with prompt optimization via BetterTogether achieves the strongest overall performance in most evaluated settings.The authors characterize the relationship between weight and prompt optimization as complementary.Source: §7 · p. 9
• The study is limited to 8-billion-parameter models, LoRA fine-tuning, two mmGRPO formulations, and a limited-feedback Banking77 setting.The authors also report that reward-only training does not yet achieve strong Banking77 performance in their setup.Source: §7 · p. 9

A Forming module-level groups

• The grouping algorithm computes trajectory rewards, records each module trace by relative invocation index, then applies padding and diversity-based selection to produce fixed-size groups.The algorithm uses PadGroups and SelectKDiverseElements after collecting traces.Source: A Forming module-level groups · p. 12
• FormModuleLevelGroups builds groups from module-level input, output, and final trajectory-reward triples aligned by module identity and relative invocation position.It returns the resulting groups together with the corresponding LM weight references.Source: A Forming module-level groups · p. 12
• When module invocation counts vary across trajectories, PadGroups either truncates incomplete invocation positions or fills them by sampling existing invocations, with fill used in the experiments.Source: A Forming module-level groups · pp. 12–13
• SelectKDiverseElements adjusts groups to the target size while favoring selections that increase reward variance.It supports both down-sampling and up-sampling.Source: A Forming module-level groups · p. 13

• DSPy implementations for the Banking77, PAPILLON, and HoVer LM programs are presented, with example datapoints supplied elsewhere in the paper.Source: B Task details · p. 13

• A DSPy program for Banking77 is provided in Python Snippet 1.Source: §B.1 · p. 13

• PAPILLON is implemented as a DSPy program in Python Snippet 2.Source: §B.2 · p. 13

• A complete DSPy program for HoVer is provided as Python Snippet 3.Source: §B.3 · pp. 14–15
• HoVer generates up to three follow-up queries to find essential clues for verifying or refuting a claim.Fewer queries may be emitted when sufficient evidence is available.Source: §B.3 · p. 14
• Each hop extracts new learnings from search results and extends the key facts used to assess the claim.The process aims to collect facts needed to identify all documents implicated by the claim.Source: §B.3 · p. 14
• The forward pass iterates over hops, searches with generated queries, ranks results by score, and retains unique documents.The final hop can use a distinct query budget and document-retention count.Source: §B.3 · p. 14
• The implementation removes duplicate documents while appending search results to the unique-document collection.Source: §B.3 · p. 14

C Example DSPy traces demonstrating disjoint module contexts

• DSPy gives each module its own locally constructed prompt and disjoint context during multi-module execution.The trace uses a simplified version of the HoVer program.Source: C Example DSPy traces demonstrating… · p. 15
• That module outputs reasoning and new key facts within the displayed interaction structure.Source: C Example DSPy traces demonstrating… · p. 15
• The HoVer key-fact extraction module receives a claim, existing key facts, and new search results as inputs.Source: C Example DSPy traces demonstrating… · p. 15

D Asset information

• The models and datasets are accessed through HuggingFace, with license information documented for the assets used.Source: D Asset information · p. 17
• Qwen3-8B uses Apache 2.0, Llama 3.1 8B Instruct uses the Meta Llama 3 Community License, and Banking77 and HoVer use CC BY 4.0.Source: D Asset information · p. 17