Search Beyond What Can Be Taught: Evolving the Knowledge Boundary in Agentic Visual Generation

Technical Breakdown§

Core Problem§

Visual generators trained on fixed corpora suffer from a world-knowledge bottleneck: they cannot handle prompts about unseen entities, trending topics, or post-cutoff events. The paper introduces a teach-then-search co-training framework to iteratively evolve the knowledge boundary—the dividing line between what a generator can internalize and what must be fetched externally.

Methodology§

Knowledge Boundary Discovery

The knowledge boundary is discovered via a teach-then-search loop:

1. Teach: Fine-tune the generator on a small set of failure cases to expand its internal knowledge.
2. Search: Use a retrieval-augmented generation (RAG) pipeline to supply external context for prompts that the generator still fails on.
3. Evaluate: The generator's performance on a held-out set guides the next iteration.

Formally, let $G_\theta$ be a generator with parameters $\theta$. For a prompt $x$, the generator produces an image $y = G_\theta(x)$. A failure classifier $F$ labels whether $y$ is correct. The knowledge boundary is the set of prompts where $F(G_\theta(x)) = \text{fail}$. The co-training minimizes:

$$ \mathcal{L}(\theta) = \mathbb{E}_{x \in \mathcal{D}_{\text{train}}} \left[ \ell(G_\theta(x), y_{\text{target}}) \right] + \lambda \cdot \mathbb{E}_{x \in \mathcal{D}_{\text{fail}}} \left[ \ell(G_\theta(x \oplus r), y_{\text{target}}) \right] $$

where $\mathcal{D}_{\text{train}}$ is the original training set, $\mathcal{D}_{\text{fail}}$ is a set of failure prompts augmented with retrieved context $r$, and $\oplus$ denotes concatenation of the prompt and context.

Key Components§

• SearchGen-20K / SearchGen-Bench: 20,839 prompts across 12 failure categories (e.g., new entities, trending events) and 22 domains.
• SearchGen-Corpus-1M: Pre-executed multimodal search corpus for offline experimentation.
• Co-training: Alternates between fine-tuning (teach) and retrieval augmentation (search).

Implementation Detail§

A minimal co-training loop:

from transformers import AutoModelForCausalLM, AutoTokenizer
from retrieval import retrieve_context

model = AutoModelForCausalLM.from_pretrained("visual-gen-model")
tokenizer = AutoTokenizer.from_pretrained("visual-gen-model")
failure_detector = FailureDetector()  # pre-trained classifier

for epoch in range(num_epochs):
    for batch in train_loader:
        prompts = batch["prompt"]
        images = batch["image"]
        
        # Phase 1: Teach (standard training)
        outputs = model(prompts)
        loss = compute_loss(outputs, images)
        loss.backward()
        
        # Phase 2: Search (for failure prompts)
        with torch.no_grad():
            generated = model.generate(prompts)
            fail_mask = failure_detector(generated, prompts)
        failed_prompts = [p for p, f in zip(prompts, fail_mask) if f]
        contexts = retrieve_context(failed_prompts, corpus=search_corpus)
        
        # Augment prompts with context and train again
        augmented_prompts = [p + " " + c for p, c in zip(failed_prompts, contexts)]
        aug_outputs = model(augmented_prompts)
        aug_loss = compute_loss(aug_outputs, [images[i] for i in fail_mask])
        aug_loss.backward()
        optimizer.step()
        optimizer.zero_grad()

Experimental Insights§

• Naive search degrades performance: generator-internal prompts suffer from noise.
• Co-training yields monotonic improvement, from a baseline of 21-28/100 to above 60/100 on SearchGen-Bench.
• The framework enables recursive self-improvement as the knowledge boundary evolves.