ComBEE parallel scan aggregation for scalable prompt learning

[nuglifeleoji]

Hi, this is a pr that implements this paper's algorithm on gepa, here are the overall summary and usage, feel free to ask questions and comment!

Summary

Adds aggregation_method parameter to support ComBEE (Combee: Scaling Prompt Learning for Self-Improving Language Model Agents), a Map-Shuffle-Reduce framework that prevents context overload when using large reflection batch sizes.

When aggregation_method="naive" (default), the existing code path is completely unchanged. When set to "combee", the reflection step applies:

• Augmented shuffling (§3.2) — each reflection is duplicated p=2 times and the full set is shuffled, ensuring every reflection has multiple chances to be incorporated
• Parallel scan aggregation (§3.1) — reflections are split into k = ⌊√n⌋ groups; Level-1 calls one reflection LM per group to produce k intermediate instructions; Level-2 aggregates those k instructions into a single final context update

This resolves the context overload problem where naive large-batch reflection degrades quality because the aggregator LLM defaults to retaining only generic patterns while discarding specific, high-value insights.

Usage

# gepa.optimize
result = gepa.optimize(
    ...,
    reflection_minibatch_size=40,   # this is n in the paper
    aggregation_method="combee",    # k = ⌊√40⌋ = 6 groups computed automatically
)

# optimize_anything
config = GEPAConfig(
    reflection=ReflectionConfig(
        reflection_minibatch_size=40,
        aggregation_method="combee",
    )
)

aggregation_method="naive" is the default — no behaviour change for existing users.

Implementation

• src/gepa/strategies/instruction_proposal.py — added default_aggregation_prompt_template to InstructionProposalSignature, the Level-2 reduce prompt that synthesizes k intermediate instructions into one
• src/gepa/proposer/reflective_mutation/reflective_mutation.py — added _propose_new_texts_combee() implementing the full Map-Shuffle-Reduce logic; propose() routes to it when aggregation_method="combee" and reflection_lm is set (falls back to naive otherwise)
• src/gepa/api.py — aggregation_method exposed on gepa.optimize()
• src/gepa/optimize_anything.py — aggregation_method added to ReflectionConfig with docstring

ComBEE is only active when reflection_lm is provided and no custom adapter proposer overrides the default LM-based reflection.

[semanticdiff-com]

Review changes with  

Changed Files

File	Status
src/gepa/proposer/reflective_mutation/reflective_mutation.py	4% smaller
docs/docs/guides/combee.md	Unsupported file format
docs/mkdocs.yml	0% smaller
src/gepa/api.py	0% smaller
src/gepa/optimize_anything.py	0% smaller
src/gepa/strategies/instruction_proposal.py	0% smaller

[nuglifeleoji]

Are there any other problems, I will be happy to fix that! @LakshyAAAgrawal Thanks for your response!

[Shangyint]

This is a nice addition, with the following caveats:

1. Can you add docs too, so this feature is not buried in the code?
2. In the feature flag comment (i'd suggest change aggregation_method to use_combee, can you also either points to the docs or the paper? Naively, gepa's minibatch size is 3, so k=sqrt(3) does not make too much sense here. Ideally, people who want to use combee should know this and set minibatch to something larger.

[Shangyint]

LGTM

[LakshyAAAgrawal]

Can you resolve the conflicts?

[nuglifeleoji]

I think all conflicts are resolved now! @LakshyAAAgrawal

[LakshyAAAgrawal]

I think all of these should go into both api.py and optimize_anything. You can add another level of config in GEPAConfig in optimize_anything called ComBEEConfig inside reflective mutation config?

[nuglifeleoji]

Just fixed all the conflicts! @LakshyAAAgrawal

[nuglifeleoji]

Thanks for this design @LakshyAAAgrawal, do I wait until Phase 1 and Phase 2 are completed?

[khram2003]

Hello! Our team at JetBrains Research is currently running experiments with GEPA, and ComBEE-style parallelization would significantly speed them up.

Could you please share when the relevant PRs are expected to be merged and ready for use?

Thanks in advance!