ComBEE: Scalable Parallel Prompt Learning¶

ComBEE (arXiv:2604.04247) scales GEPA's reflection step to large minibatches without quality degradation. It ships as a ReflectionLM implementation: pass reflection_strategy=ComBEEReflectionLM(...) to gepa.optimize, or set ReflectionConfig(reflection_strategy=...) in optimize_anything.

How ComBEE works: Map-Shuffle-Reduce¶

ComBEE replaces the single reflection call with a three-phase pipeline, per component:

n traces → [Augmented Shuffle] → k groups → [Level-1: k LM calls] → k proposals → [Level-2: 1 LM call] → final instruction

1. Augmented Shuffle (§3.2)¶

Each reflection record is duplicated p times (duplication_factor, default 2) and the augmented set is shuffled with a seeded RNG before being distributed across groups. Every record gets multiple chances to be incorporated, improving robustness at large batch sizes — and runs stay reproducible.

2. Level-1 — Map (§3.1)¶

The augmented set (p·n items) is split into k = ⌊√n⌋ groups. One reflection-LM call per group produces k intermediate instruction proposals. Each group sees p·n/k ≈ p·√n traces — a manageable context even when n is large. With n=40: k = 6 groups, each seeing ~13 traces instead of 40.

3. Level-2 — Reduce (§3.1)¶

A final LM call synthesizes the k intermediate proposals into one instruction (aggregation_prompt_template, customizable). The choice k = ⌊√n⌋ balances both levels: Level-1 processes √n traces per group, Level-2 aggregates √n proposals — both at the same scale.

Usage¶

Raise reflection_minibatch_size

GEPA's default reflection_minibatch_size is 3. With n < 4, k = ⌊√n⌋ = 1 and ComBEE falls back to a standard single reflection call (it logs when this happens). Set reflection_minibatch_size to 20 or more for meaningful benefit.

`gepa.optimize`¶

import gepa
from gepa.proposer.reflective_mutation.combee import ComBEEReflectionLM

result = gepa.optimize(
    seed_candidate={"system_prompt": "You are a helpful assistant."},
    trainset=trainset,
    valset=valset,
    task_lm="openai/gpt-4.1-mini",
    reflection_strategy=ComBEEReflectionLM("openai/gpt-5.1"),
    reflection_minibatch_size=40,  # n — ComBEE forms k=6 groups automatically
    max_metric_calls=600,
)

`optimize_anything`¶

Through the engine-pluggable API, engine_config maps onto GEPAConfig field-for-field:

from gepa.optimize_anything import OptimizeAnythingConfig, ReflectionConfig, optimize_anything
from gepa.proposer.reflective_mutation.combee import ComBEEReflectionLM

result = optimize_anything(
    seed_candidate=seed,
    evaluator=my_evaluator,
    dataset=dataset,
    config=OptimizeAnythingConfig(
        engine="gepa",
        max_evals=600,
        engine_config={
            "reflection": ReflectionConfig(
                reflection_strategy=ComBEEReflectionLM("openai/gpt-5.1"),
                reflection_minibatch_size=40,
            ),
        },
    ),
)

The legacy launcher config is also still accepted directly:

from gepa.optimize_anything import GEPAConfig, ReflectionConfig, optimize_anything
from gepa.proposer.reflective_mutation.combee import ComBEEReflectionLM

result = optimize_anything(
    seed_candidate=seed,
    evaluator=my_evaluator,
    dataset=dataset,
    config=GEPAConfig(
        reflection=ReflectionConfig(
            reflection_strategy=ComBEEReflectionLM("openai/gpt-5.1"),
            reflection_minibatch_size=40,
        ),
    ),
)

Options¶

ComBEEReflectionLM(
    lm,                                # model name string, or any LanguageModel callable
    lm_kwargs=None,                    # completion options for a model-name lm
    reflection_prompt_template=None,   # Level-1 template: str, or dict per component
    aggregation_prompt_template=None,  # Level-2 template (must contain <curr_param> and <side_info>)
    duplication_factor=2,              # p — augmented-shuffle duplication (§3.2)
    rng=None,                          # None -> engine-bound RNG; pass random.Random(...) for an independent stream
    logger=None,                       # GEPA injects its configured logger when omitted
    batch_reflection=True,             # batches only when the LM provides batch_complete
)

Seeding the shuffle

By default GEPA binds ComBEE to the engine RNG. This preserves the single-proposal behavior of #307: shuffles participate in the same random stream as candidate selection and minibatch sampling. Same-seed runs therefore reproduce the legacy call and shuffle sequence. Passing an explicit rng=random.Random(...) opts into an independent, pinned shuffle stream instead.

Public reflection configuration

When ComBEE is supplied as reflection_strategy=, GEPA forwards the public reflection_prompt_template and reflection_lm_kwargs settings to it. Constructor values on ComBEEReflectionLM take precedence. Plain callable LMs provide token estimates only; use a cost-tracking LM when setting max_reflection_cost.

Cost and observability¶

ComBEE makes k + 1 reflection-LM calls per component per proposal (vs 1 for the default reflector). The per-call intermediates are recorded in proposal metadata under combee:-namespaced keys — combee:<comp>:k, combee:<comp>:level1_prompts / level1_outputs, combee:<comp>:num_lm_calls, combee:<comp>:mode, and combee:total_lm_calls — visible to on_proposal_end consumers (in the event's metadata) and experiment trackers (in the proposal_reflection_metadata table). total_cost and token totals are exposed by delegation to the wrapped LM, so max_reflection_cost works as a stop condition.

When the LM provides batch_complete, ComBEE issues the iteration's Level-1 map calls as one batched wave and the Level-2 reduce calls as a second. That includes the default single-proposal path: one job's k maps go out together, and a singleton reduce uses a plain completion (a wave of one is not batched). Under parallel proposals, maps from all proposals share the first wave and reduces share the second. This assumes the LM's calls are exchangeable (each reply depends only on its own prompt — the standard property of stateless completion APIs). Without that capability, ComBEE automatically executes complete jobs in strict #307 order, preserving sequential results for ordinary and order-dependent callables. Set batch_reflection=False to request that strict ordering even for a batch-capable LM. Failed attempts restore the RNG state and memoize already-completed logical calls, so both reflect() / reflect_many retries and the engine's per-job recovery path preserve results without repurchasing work.

Credits¶

ComBEE support was originally contributed by @nuglifeleoji in #307 and re-hosted onto the ReflectionLM protocol introduced in #369.