# How Ablation Studies Are Performed in Stage 4 of AI Scientist v2

> Discover how SakanaAI's AI Scientist v2 performs ablation studies in Stage 4. Learn how the ParallelAgent class isolates component impact and refines AI models.

- Repository: [Sakana AI/AI-Scientist-v2](https://github.com/SakanaAI/AI-Scientist-v2)
- Tags: how-to-guide
- Published: 2026-03-28

---

**Ablation studies in Stage 4 of AI Scientist v2 are orchestrated by the `ParallelAgent` class, which systematically generates, executes, and tracks experiments that remove or modify components from the best Stage 3 implementation to isolate their impact.**

Stage 4 serves as the experimental validation phase of the [SakanaAI/AI-Scientist-v2](https://github.com/SakanaAI/AI-Scientist-v2) pipeline. Unlike earlier stages that focus on ideation and implementation, this stage rigorously tests the contribution of individual architectural choices using an LLM-guided workflow that prevents duplicate experiments and maintains scientific rigor.

## Initialization of the ParallelAgent

When the pipeline transitions to Stage 4, the `AgentManager._create_agent_for_stage` method in [`ai_scientist/treesearch/agent_manager.py`](https://github.com/SakanaAI/AI-Scientist-v2/blob/main/ai_scientist/treesearch/agent_manager.py) constructs a dedicated `ParallelAgent` instance. This agent receives the **best implementation from Stage 3** via the `best_stage3_node` parameter, establishing the baseline code that will undergo systematic ablation.

The initialization logic specifically checks for `main_stage == 4` to trigger this specialized agent configuration:

```python

# Inside AgentManager._create_agent_for_stage

if main_stage == 4:
    # Retrieve the best node from Stage 3

    best_stage3_node = self._get_best_implementation(last_substage.name)

agent = ParallelAgent(
    task_desc=task_desc,
    cfg=stage_cfg,
    journal=self.journals[stage.name],
    stage_name=stage.name,                     # e.g. "4_ablation"

    best_stage3_node=best_stage3_node,        # baseline code to ablate

)

```

## Baseline Selection and Node Handling

The `ParallelAgent._select_parallel_nodes` method detects Stage 4 execution through a simple prefix check on `self.stage_name`. When the stage name starts with `"4_"`, the agent **automatically adds `self.best_stage3_node` to the processing queue** according to lines 2006-2012 in [`ai_scientist/treesearch/parallel_agent.py`](https://github.com/SakanaAI/AI-Scientist-v2/blob/main/ai_scientist/treesearch/parallel_agent.py).

This design ensures that every parallel worker operates on the identical baseline implementation, preventing the redundant ablation attempts that would occur if workers selected different starting points.

## Generating Ablation Ideas with LLM Prompting

The core of Stage 4's scientific methodology resides in `_generate_ablation_idea`, a private method that constructs targeted prompts for the LLM. This method builds a structured context containing:

- The current implementation code (`self.best_stage3_node.code`)
- A history of completed ablations (`self._ablation_state["completed_ablations"]`)
- Specific instructions to propose **one new, distinct ablation** (such as disabling a specific component or swapping a dataset)

The prompt engineering explicitly requires the LLM to identify a single component to ablate and ensure it differs from previous attempts:

```python
def _generate_ablation_idea(self) -> Optional[AblationIdea]:
    completed = list(self._ablation_state["completed_ablations"])
    prompt = {
        "Introduction": "You are an AI researcher conducting ablation studies...",
        "Base code you are working on": wrap_code(self.best_stage3_node.code),
        "Previous Ablations": {"Has been tried": completed or "Nothing has been tried yet."},
        "Instructions": {"Requirements": [
            "1. Identify ONE specific component/feature to ablate",
            "2. Ensure the ablation is different from previous attempts",
            "3. ..."]},
        "Response format": "ABLATION NAME: ...\\nABLATION DESCRIPTION: ..."
    }
    response = query(system_message=prompt,
                     model=self.cfg.agent.code.model,
                     temperature=self.cfg.agent.code.temp)
    name, desc = _parse_keyword_prefix_response(
        response, "ABLATION NAME:", "ABLATION DESCRIPTION:")
    return AblationIdea(name=name, description=desc) if name and desc else None

```

If parsing fails after configured retries, the system implements a **fallback mechanism** that returns a default ablation idea (adding a layer) to maintain workflow continuity.

## Tracking Ablation State and Preventing Duplicates

Stage 4 maintains rigorous experiment tracking through the `_ablation_state` dictionary, which uses a `"completed_ablations"` set to ensure scientific validity. The workflow implements **immediate reservation** of ablation ideas:

When a new ablation is generated, it is immediately added to `self._ablation_state["completed_ablations"]` (lines 2111-2113) before execution begins. This prevents other parallel workers from proposing identical experiments even while the first worker is still running the code.

After execution completes, `_update_ablation_state` validates the result:

```python
def _update_ablation_state(self, result_node: Node):
    if not self.stage_name or not self.stage_name.startswith("4_"):
        return
    if result_node.ablation_name and not result_node.is_buggy:
        self._ablation_state["completed_ablations"].add(result_node.ablation_name)
        logger.info(f"Ablation {result_node.ablation_name} completed successfully")

```

Successful runs (non-buggy nodes) are recorded permanently, while failed experiments are logged but excluded from the completion set, allowing the system to retry or propose alternatives.

## Parallel Execution Architecture

The `ParallelAgent` spawns a process pool (`self.executor`) that maps workers across available GPUs or CPU cores. Each worker receives:
- The baseline node data from Stage 3
- The generated ablation idea (name and description)
- Auxiliary context such as plot code from prior stages

Workers execute the modified code, capture performance metrics, and return result nodes to the main process, which inserts them into the stage journal for analysis.

## Completion Criteria and Configuration

Stage 4 terminates based on criteria defined in [`bfts_config.yaml`](https://github.com/SakanaAI/AI-Scientist-v2/blob/main/bfts_config.yaml). The workflow respects the **`stage4_max_iters: 18`** limit by default, or stops when no further improvements are observed against the evaluation metric defined in `_define_global_metrics`.

Upon completion, the system aggregates results through [`perform_experiments_bfts_with_agentmanager.py`](https://github.com/SakanaAI/AI-Scientist-v2/blob/main/perform_experiments_bfts_with_agentmanager.py), incorporating ablation summaries into the final research report.

## Summary

- **ParallelAgent orchestration**: Stage 4 uses a specialized agent that receives the best Stage 3 implementation via `best_stage3_node` in [`agent_manager.py`](https://github.com/SakanaAI/AI-Scientist-v2/blob/main/agent_manager.py).
- **Systematic idea generation**: The `_generate_ablation_idea` method in [`parallel_agent.py`](https://github.com/SakanaAI/AI-Scientist-v2/blob/main/parallel_agent.py) uses structured LLM prompting to propose unique component removals.
- **Duplicate prevention**: The `_ablation_state` dictionary tracks `completed_ablations` immediately upon idea generation and confirms them after successful execution.
- **Parallel execution**: Process pools distribute ablation experiments across hardware resources while maintaining shared state.
- **Configured limits**: Execution respects `stage4_max_iters` from [`bfts_config.yaml`](https://github.com/SakanaAI/AI-Scientist-v2/blob/main/bfts_config.yaml) and global metric thresholds.

## Frequently Asked Questions

### What is the purpose of Stage 4 in AI Scientist v2?

Stage 4 conducts **ablation studies** to isolate and measure the contribution of specific architectural components or features within the best implementation discovered in Stage 3. This stage transforms promising implementations into rigorously validated scientific findings by systematically removing elements and measuring performance deltas.

### How does the system prevent running the same ablation twice?

The system implements a **reservation pattern** through `_ablation_state["completed_ablations"]`. When an ablation idea is generated, it is immediately added to this set before execution begins. The `_update_ablation_state` method only permanently records successful runs, while failed attempts remain available for retry or replacement with alternative ablations.

### What happens if the LLM fails to generate a valid ablation idea?

The `_generate_ablation_idea` method includes a **fallback mechanism** (lines 1910-1919 in [`parallel_agent.py`](https://github.com/SakanaAI/AI-Scientist-v2/blob/main/parallel_agent.py)) that activates after multiple parsing retries fail. When the LLM response cannot be parsed for "ABLATION NAME" and "ABLATION DESCRIPTION", the system defaults to a generic fallback idea—typically adding a layer—to ensure the workflow continues rather than stalling.

### Where is the baseline code for ablation studies stored?

The baseline code originates from **Stage 3's best implementation**, passed through the `best_stage3_node` parameter during `ParallelAgent` initialization in [`agent_manager.py`](https://github.com/SakanaAI/AI-Scientist-v2/blob/main/agent_manager.py). This node contains the complete code that workers modify according to each ablation specification, ensuring all experiments share a consistent starting point.