Source-linked AI summary
Context Training with Active Information Seeking
Zeyu Huang, Adhiguna Kuncoro, Qixuan Feng, Jiajun Shen, Lucio Dery, Arthur Szlam, Marc'Aurelio Ranzato
TL;DR
Deployed LLMs are difficult and costly to adapt when tasks require new or niche knowledge, while existing context optimizers generally rely on closed-loop internal knowledge. The paper adds Wikipedia search and browser tools, then uses search-based training to maintain and prune candidate contexts. Across translation, health, and reasoning benchmarks, this combination produces consistent improvements, with data efficiency and cross-model transfer, while depending on the executor model’s ability to use the constructed context.
Problem
Existing context optimization methods remain closed-loop, limiting their ability to acquire task-relevant knowledge absent from the model’s parametric memory.
Method
The paper equips context optimizers with Wikipedia search and browser tools and trains them with beam-search-style maintenance and pruning of multiple candidate contexts.
Results
Across low-resource translation, healthcare, and reasoning benchmarks, active information seeking with search-based training yields consistent improvements over passive or sequential baselines.
Takeaways & Limitations
The approach is data-efficient, robust across hyperparameters, and transfers effective textual contexts across models.
Takeaways & Limitations
The method relies on the base model’s context-utilization ability, and the constructed context is often sparse, pointwise, and data-specific.
Abstract
from arXiv · showhide
Most existing large language models (LLMs) are expensive to adapt after deployment, especially when a task requires newly produced information or niche domain knowledge. Recent work has shown that, by manipulating and optimizing their context, LLMs can be tailored to downstream tasks without updating their weights. However, most existing methods remain closed-loop, relying solely on the model's intrinsic knowledge. In this paper, we equip these context optimizers with Wikipedia search and browser tools for active information seeking. We show that naively adding these tools to a standard sequential context optimization pipeline can actually degrade performance compared to baselines. However, when paired with a search-based training procedure that maintains and prunes multiple candidate contexts, active information seeking delivers consistent and substantial gains. We demonstrate these improvements across diverse domains, including low-resource translation (Flores+), health scenarios (HealthBench), and reasoning-heavy tasks (LiveCodeBench and Humanity's Last Exam). Furthermore, our method proves to be data-efficient, robust across different hyperparameters, and capable of generating effective textual contexts that generalize well across different models.
1. Introduction
Context optimization adapts deployed LLMs without weight updates, but closed optimizers cannot reliably acquire missing external knowledge. The paper addresses this with active information seeking and beam-search-style training that controls noisy or suboptimal updates.
- Motivation: Context optimization adapts LLMs by refining their input context or working memory instead of updating model parameters.This avoids retraining or fine-tuning costs while supporting task specialization.
- Problem: Closed optimizers struggle when task-relevant information lies outside the model’s parametric memory or feedback does not contain the missing knowledge.They can mainly rearrange or extrapolate from existing knowledge, limiting effective context updates.
- Failure modes: Web access introduces context pollution from misleading information and local optima from greedy, sequential optimization.Without explicit backtracking, optimizers may fail to recover from harmful updates or miss better trajectories.
- Approach: Beam-search-style training maintains multiple candidate contexts, explores updates in parallel, and prunes noisy or weak trajectories.The candidate pool includes a “Do Nothing” option that retains the current best context when new explorations are unhelpful.
- Findings: Active information seeking paired with search-based training delivers consistent gains over sequential, closed-context training without task-specific prompt tuning.The method is also reported as data-efficient, hyperparameter-robust, and transferable across models.
2. Related Work: from Context Engineering to Working-Memory Evolution
Prior work evolves context from static input composition toward dynamic working memory, but its optimization is typically closed and dependent on internal knowledge. This paper studies externally grounded information seeking during context optimization using a general framework shared across domains.
- Self-Evolving Working-Memory: Working-memory methods transform static context into a dynamic workspace for adapting LLM agents to unfamiliar tasks and environments.This line of work extends in-context learning beyond one-time input composition.
- Context Engineering: Context engineering optimizes the information distribution in a frozen LLM’s input, including instructions, retrieved documents, and tool outputs.Its goal is to construct a composite input that maximizes inference capability.
- Open Gap: Existing context optimization typically operates as a closed system that relies on environmental feedback and the optimizer’s internal reflection capabilities.This raises whether an optimizer can update context effectively when it lacks prerequisite knowledge.
- This Work: The paper equips the optimizer with information-seeking capabilities to retrieve external information during context optimization.Unlike standard retrieval-augmented generation, the optimizer actively seeks missing information rather than assuming an existing corpus supplies it.
- Implication: A shared, non-task-specific framework shows that external grounding becomes effective with search-based training and can integrate with many existing agent workflows.The authors contrast this with the continued constraints of standard sequential training under frozen model knowledge.
3. Methodology
The methodology treats learning as optimization over a discrete, human-readable context while keeping model weights frozen. It adds structured context management, external information-seeking tools, and beam-search training to mitigate pollution and local-optimum failures in sequential optimization.
- Learning as State Optimization: The framework optimizes a modifiable state using predictions, feedback, an optimizer, task distribution, and reward.At each step, the optimizer updates the current state from a learnable batch to reduce prediction discrepancy.
- Frozen-Weight Context Training: Context training changes model behavior without altering weights, using executor and optimizer agents in an optimization-like pipeline.The executor solves tasks using the current context, while the optimizer reads trajectories and feedback to update it.
- Context Management: The context is implemented as a structured database whose resource items can be read, written, updated, deleted, and searched through tool calls.Resources include identifiers, summaries, raw content, metadata, keywords, and embeddings.
- Information Seeking: Wikipedia search and browser tools give the optimizer external grounding for declarative gaps and more complex or recently indexed information.These tools allow the optimizer to retrieve missing information beyond its frozen parametric knowledge.
- Sequential Training Pitfalls: Sequential training is vulnerable to context pollution: a roughly 200-token update was associated with a precipitous performance decline, followed by prolonged failure to recover.The optimizer repeatedly added and removed information rather than pruning the harmful content, motivating explicit backtracking.
- Sequential Training Pitfalls: Sequential optimization can also remain in a local strategy basin, repeatedly pruning and re-adding dictionary resources instead of discovering stronger alternatives.The observed context length shows a sawtooth pattern, while the optimizer fails to shift toward strategies such as increasing Parallel Examples.
- Beam-Search Training: Beam-search training maintains K candidate contexts, expands each with alternative updates, and prunes candidates using held-out validation feedback.The candidate pool includes the previous best context as a “Do Nothing” option, preserving it when new explorations are noisy or harmful.
4. Experiments
The experiments evaluate active information seeking across translation, health, and complex reasoning tasks, comparing sequential and beam-search context optimization with and without external tools. BeamSearch-IS consistently outperforms relevant baselines, while naïvely adding information seeking to sequential training can reduce performance.
- Evaluation setup: The evaluation spans low-resource translation, HealthBench, LiveCodeBench, and Humanity’s Last Exam using held-out test metrics.The study compares ChrF++, rubric-based health scores, pass@1/pass@8, and average@8 accuracy under constrained training data.
- Low-resource translation: Seq-IS lowers average translation performance from 31.13 for Seq to 29.68, illustrating the cost of unverified information seeking.This result is reported across five English-to-low-resource-language pairs.
- Low-resource translation: BeamSearch-IS reaches an average translation score of 34.51, improving 4.14 points over BoN and exceeding Gemini-2.5-Pro at 30.37.The result indicates that beam-search context selection can distinguish useful evidence from noisy retrievals.
- HealthBench: On HealthBench, Seq-IS scores 0.4484 versus 0.4629 for Seq, whereas BeamSearch-IS reaches 0.5026, comparable to Gemini-2.5-Pro at 0.5030.BeamSearch-IS performs especially well on Health Data Tasks and Emergency Referrals, while Gemini-2.5-Pro leads in Response Depth.
- Complex reasoning: BeamSearch-IS improves LCB Hard pass@1 to 33.9% versus a 30.0% baseline and achieves 8.63% average accuracy on HLE.Internal-only methods yield negligible gains, while Seq-IS is inconsistent across LCB and HLE.
5. Analysis
The analyses examine how BeamSearch-IS explores contexts, uses limited data, tolerates hyperparameter variation, and transfers across models. They show that exploration selects better resource combinations, reaches strong translation performance with few samples, and produces broadly transferable contexts.
- Exploration: Beam Search moves from Dictionary Support toward Linguistic Rules combined with Parallel Examples, a trajectory associated with better translation quality.The analysis frames this movement as exploration that avoids local optima rather than remaining with the sequential baseline’s less effective strategy.
- Data efficiency: BeamSearch-IS reaches a score above 23.0 with only 32 training samples on English-to-Southwestern Dinka.Standard methods remain in lower-performance regions even as the data volume increases.
- Hyperparameter robustness: BeamSearch-IS maintains scores around 22.2–22.45 across diverse balanced hyperparameter configurations.Performance drops to 20.73 under the extreme 6-1-1 configuration, indicating robustness within reasonable parameter balances.
- Data utility: Utility maps are sparse and modular, but continuous vertical bands identify dominant resources that benefit many test examples.The resource axis is arranged from longer to shorter contexts, and the test-data axis indexes individual examples.
- Generalization: Contexts optimized on Gemini-2.5-Flash transfer to Gemini-3-Flash, with BeamSearch-IS improving performance across all evaluated domains.The transferred context raises Magahi to 52.12 and HealthBench to 0.6624, whereas closed Seq often transfers poorly.
6. Conclusion
The paper combines active information seeking with search-based context training to overcome the limitations of passive, closed context optimization. Across low-resource translation, healthcare, and reasoning benchmarks, this combination consistently improves performance and transfers across models.
- The method augments context optimizers with Wikipedia search and browser-based tools for active information seeking.
- Naively adding search tools to sequential context training can be harmful, whereas search-based training makes them effective in practice.
- Across low-resource translation, healthcare, and reasoning benchmarks, the approach yields consistent improvements over passive baselines.
- The approach remains data-efficient and transfers effectively across different models.
7. Limitations and Future Work
The method's effectiveness depends on the base model's ability to use retrieved context, while sparse and highly instance-specific resources can limit generalization. Future work targets stronger context utilization, broader search, and hybrid offline-online settings.
- Limitations: Performance depends on the base model's ability to utilize the retrieved context effectively.The authors report larger HLE gains when context trained with Gemini-2.5-Flash is applied to the more capable Gemini-3-Flash model.
- Limitations: The constructed context is sparse, with most collected resources being pointwise and highly data-specific, which can hinder generalization to diverse test instances.
- Future Work: Future work includes improving context utilization, enabling broader exploration, and combining offline background preparation with online exploration.
8. Appendix
The appendix specifies the beam-search context-training algorithm, benchmark data construction, training settings, and implemented tooling. The procedure initializes and validates contexts, expands candidate trajectories, retains the best state, and uses task-specific data splits and configurations.
- Algorithm: The algorithm takes training and validation data, beam width K, branching factor M, and optimization steps per child L as inputs.
- Algorithm: Training begins with an empty context, validates it, and initializes the beam and global best context.
- Algorithm: Each global step expands and optimizes candidate contexts while retaining the current global best as a candidate.
- Algorithm: Selection uses elitism to update the global best context from the candidate pool.
- Data and Training: FLORES++, HealthBench, and Humanity's Last Exam use specified training, validation, and test subsets, with HLE sampling designed around instance-specific knowledge.HLE problems are grouped into eight high-level categories, and an LLM scores within-category relevance to make the training set representative.
- Data and Training: BeamSearch-IS and BeamSearch use K=2 and M=3 across tasks, with epoch counts differing between translation and other tasks.Seq-IS and Seq use separate epoch schedules for fair comparison.
### ROLE & MISSION
The optimizer is framed as an autonomous context engineer that improves structured context for a task. Its output is delivered to a downstream executor to increase success on similar tasks.
- ROLE & MISSION: The optimizer agent performs one-step optimization of structured context for a given task.
- ROLE & MISSION: The optimized context is delivered to a downstream executor agent for task performance on similar tasks.
### TOOL AUGMENTATION
The system augments the context-management process with dynamically specified tools, whose descriptions, inputs, and outputs are supplied to the agent.
- The agent receives a dynamically populated list of tools for managing context.Each tool is accompanied by its name and description.
### TOOL USAGE RULES
The tool-usage rules require tool calls to follow the correct argument format and prohibit repeating identical calls.
- Every response must call a tool and finish by calling final_answer_tool.
- Tool calls must use correct arguments and cannot repeat the exact same tool with identical parameters.
### INCOMING DATA PACKAGE EXPLANATION
The incoming package supplies context-management state, executor feedback, task prompts, tool instructions, and optimization guidance for improving context generalization.
- Incoming data: The package includes context-update history, a context preview, and executor trajectories with task and output information.It may also include reference answers, evaluation results, context-usage summaries, and previous update attempts.
- Optimization guidance: The optimizer is instructed to improve the context for future similar tasks rather than only the current task.
- Optimization guidance: The context should be understandable, with information organized so it is easy to find.
- Optimization guidance: Active searching requires using auto_browser_use_tool to seek high-quality external knowledge and trying diverse queries when needed.The workflow also recommends adding high-quality examples and relevant knowledge to support generalization.
- Optimization guidance: The optimizer may modify content but must not perform branch-management actions such as creating, checking out, merging, or committing branches.Changes are committed automatically after completion.
- Executor workflow: The executor reviews the task and context, may retrieve additional information, evaluates context quality, and submits the final answer through final_answer_tool.The context may be empty, in which case the executor relies on its internal knowledge.
- Tool specification: The package defines available tools through their names, descriptions, inputs, and output types, including dynamically populated tool sections.
- Output constraints: For final responses, the executor must use an action blob naming final_answer_tool and providing an answer argument.The instructions state that this tool is the required completion mechanism.