MOAR: Multi-Objective Agent-based Refinement

Relevant source files

The following files were used as context for generating this wiki page:

• docetl/moar/MOARSearch.py
• docetl/moar/ParetoFrontier.py
• experiments/reasoning/run_moar.py
• docetl/moar/cli_helpers.py
• experiments/reasoning/RUNMOAR.md
• docetl/reasoning_optimizer/AGENT.md
• docetl/moar/search_utils.py
• docetl/reasoning_optimizer/directives/README.md
• docetl/reasoning_optimizer/directives/deterministic_doc_compression.py
• docetl/reasoning_optimizer/directives/cascade_filtering.py

MOAR (Multi-Objective Agent-based Refinement) is an optimization system within DocETL that automates the discovery of Pareto-optimal data processing pipelines. It employs a graph-search algorithm to iteratively rewrite and evaluate pipeline configurations, balancing the trade-offs between execution cost and accuracy. MOAR uses an LLM-based agent to apply transformation rules, called Directives, to explore the space of possible pipelines.

The core of MOAR is the MOARSearch class, which implements a search algorithm inspired by Monte Carlo Tree Search (MCTS). It systematically explores a tree of pipeline configurations, where each node represents a unique pipeline YAML file, to identify a set of optimal solutions that offer the best possible accuracy for a given cost.

Sources: experiments/reasoning/RUNMOAR.md, docetl/moar/MOARSearch.py:20-26

Architecture and Core Components

The MOAR system is built around three main classes: MOARSearch, Node, and ParetoFrontier. These components work together to manage the search process, represent pipeline states, and track optimal solutions.

Sources: docetl/moar/MOARSearch.py

This diagram illustrates the relationships between the core components of the MOAR search system.

Sources: docetl/moar/MOARSearch.py, docetl/moar/ParetoFrontier.py

MOARSearch Class

The MOARSearch class is the main driver of the optimization process. It orchestrates the search algorithm, manages the state of the search tree, and interacts with other components.

Key Responsibilities:

• Initializes the search with a root pipeline configuration (root_yaml_path).
• Manages the set of available transformation Directives.
• Executes the main search loop for a specified number of iterations.
• Tracks total search cost, including LLM calls for expansion and instantiation.
• Coordinates the four main phases of the search algorithm: Selection, Expansion, Simulation, and Backpropagation.

Sources: docetl/moar/MOARSearch.py:29-93

Node Class

A Node in the search tree represents a specific pipeline configuration, defined by a YAML file. Each node stores metrics and relationships that are crucial for the search algorithm.

Key Attributes:

• yaml_file_path: The path to the pipeline configuration file.
• parent & children: Pointers to maintain the tree structure.
• cost: The execution cost of the pipeline, calculated during the simulation phase.
• value: The utility score of the node, used for selection.
• visits: The number of times the node has been visited during the search.
• latest_action: The Directive that was applied to its parent to create this node.

Sources: docetl/moar/MOARSearch.py:440-598

ParetoFrontier Class

The ParetoFrontier class is responsible for tracking the set of non-dominated solutions found during the search. A solution is on the Pareto frontier if no other solution is both cheaper and more accurate.

Key Responsibilities:

• Maintains a list of plans (frontier_plans) that are currently on the Pareto frontier.
• Stores the accuracy for each evaluated plan in plans_accuracy.
• Updates the frontier whenever a new plan is simulated via the add_plan_f1 method.
• Manages rewards and performance statistics for each Directive.

Sources: docetl/moar/ParetoFrontier.py

The MOAR Search Algorithm

MOAR uses a concurrent MCTS-like algorithm to explore the solution space. The main loop, managed by MOARSearch.search(), runs multiple search iterations in parallel using a ThreadPoolExecutor. Each iteration consists of four distinct phases.

Sources: docetl/moar/MOARSearch.py:277-319

This flowchart shows the main search loop and the four phases of a single search iteration.

Sources: docetl/moar/MOARSearch.py:277-353

1. Selection

The selection phase traverses the search tree from the root to find the most promising leaf node to expand. The select method repeatedly chooses the best child based on the Upper Confidence Bound (UCB1) formula until a leaf node (one that is not fully explored) is reached.

The best_child() method on a Node calculates the UCB1 score for each child to balance exploitation (choosing nodes with high value) and exploration (choosing nodes with few visits).

Sources: docetl/moar/MOARSearch.py:488-500, docetl/moar/search_utils.py:110-122

2. Expansion

Once a leaf node is selected, the expansion phase generates new child nodes by applying Directives. The expand method determines an optimization goal (acc for accuracy or cost) based on the node's performance relative to others. It then calls an LLM agent with a dynamically generated prompt (create_expansion_prompt_acc or create_expansion_prompt_cost) to choose a directive and target operator(s).

The agent's structured response is used to instantiate the directive, which in turn modifies the parent node's YAML to create one or more new child nodes. Some directives, like chaining, can produce multiple candidate children (multi-instance), which are all simulated before the best one is chosen.

Sources: docetl/moar/MOARSearch.py:590-843, docetl/moar/search_utils.py:141-382

The following diagram shows the interaction during the expansion phase.

3. Simulation

In the simulation phase, the simulate method evaluates each newly created child node. This involves two steps:

1. Execution: The pipeline defined by the node's YAML is executed via node.execute_plan(). This process calculates the pipeline's monetary cost.
2. Evaluation: The output of the executed pipeline is passed to a user-provided evaluation function (evaluate_func) to calculate an accuracy score.

If either step fails, the node is assigned an infinite-cost or infinite-negative-accuracy and is typically pruned from the search.

Sources: docetl/moar/MOARSearch.py:845-885

4. Backpropagation

After simulation, the results (cost and accuracy) are used to update the search tree. The backpropagate method traverses from the simulated node back to the root.

• The value of each ancestor node is updated based on whether the new node improved the Pareto frontier.
• The visits count of each ancestor is incremented.
• The ParetoFrontier updates action-specific statistics (rewards, cost changes, accuracy changes) which inform future expansion decisions.

Sources: docetl/moar/MOARSearch.py:906-920

Directives System

Directives are the core of MOAR's "Agent-based Refinement." They are modular, reusable transformation rules that an LLM agent can apply to a pipeline. Each directive encapsulates a specific optimization strategy.

Sources: docetl/reasoning_optimizer/AGENT.md, docetl/reasoning_optimizer/directives/README.md

Directive	Description	Source File
deterministic_doc_compression	Reduces LLM costs by using regex to remove irrelevant content from documents before processing.	deterministic_doc_compression.py
cascade_filtering	Injects a series of cheaper filters (code-based, then simple LLM) before an expensive filter to reduce cost.	cascade_filtering.py
doc_summarization	Adds a map operation at the start of a pipeline to summarize long documents.	doc_summarization.py
change_model	Switches the LLM model for an operation to trade off cost and accuracy.	change_model.py
gleaning	Adds a validation loop to an operation to improve output quality.	gleaning.py
chaining	Replaces a complex operation with a sequence of simpler ones.	chaining.py

Sources: docetl/reasoning_optimizer/directives/README.md, docetl/reasoning_optimizer/directives/deterministic_doc_compression.py:20-29, docetl/reasoning_optimizer/directives/cascade_filtering.py:19-28, docetl/reasoning_optimizer/directives/doc_summarization.py:20-31

Execution and Configuration

MOAR optimization is typically launched via a command-line script or a YAML configuration file.

Command-Line Execution

The experiments/reasoning/run_moar.py script provides a command-line interface for running MOAR experiments.

Parameter	Required	Description
--yaml_path	Yes	Path to the initial input YAML pipeline file.
--dataset_path	Yes	Path to the dataset file for sample input data.
--experiment_name	Yes	A unique name for the experiment run.
--dataset	Conditional	The name of the dataset for evaluation (e.g., cuad).
--max_iterations	No	Maximum number of MOARSearch iterations (default: 100).
--exploration_weight	No	UCB exploration parameter c (default: 1.414).
--model	No	The LLM model to use for directive instantiation (default: gpt-5).
--accuracy_function	No	Path to a Python file with a custom evaluation function.
--accuracy_metric_key	No	The key to extract from the evaluation results for the accuracy score.

*Sources: experiments/reasoning/run_moar.py:452-487, experiments/reasoning/RUNMOAR.md *

YAML Configuration

MOAR can also be configured and run via an optimizer_config section within a DocETL pipeline YAML file. The docetl/moar/cli_helpers.py script handles parsing this configuration.

This diagram shows the execution flow when running MOAR via the CLI helper.

Sources: docetl/moar/cli_helpers.py:84-282

Key optimizer_config parameters:

Parameter	Type	Description
save_dir	string	Directory to save all experiment outputs.
available_models	list	List of LLM names available for the change_model directive.
evaluation_file	string	Path to the Python file containing the evaluation function.
metric_key	string	The key for the primary accuracy metric in the evaluation results.
max_iterations	integer	Number of search iterations to run.
model	string	The LLM to use for the agent's reasoning.
exploration_weight	float	The UCB exploration constant c.

Sources: docetl/moar/cli_helpers.py:112-148

Outputs

A MOAR run produces several output files in the specified save directory.

• experiment_summary.json: Contains metadata about the run, including parameters, duration, and paths to other result files.
• pareto_frontier.json: A JSON array listing all solutions on the final Pareto frontier, with their associated cost and accuracy.
• evaluation_metrics.json: Detailed evaluation results for the pipelines on the Pareto frontier.
• moar_tree_log.txt: A log of the search tree structure, node visits, and action performance statistics for each iteration.
• Generated YAML files: The directory contains all the pipeline YAML files that were explored during the search, named with their node ID.

Sources: experiments/reasoning/RUNMOAR.md, docetl/moar/MOARSearch.py:961-987, docetl/moar/cli_helpers.py:253-276