acesearcher-reasoning-search-self-play

Files (1)

• SKILL.md

SKILL.md (first lines)

## AceSearcher: Bootstrapping Reasoning and Search for LLMs via Reinforced Self-Play

### Outcome

Implement a cooperative self-play framework that trains a single LLM to alternate between decomposer and solver roles, enhancing multi-hop reasoning and retrieval-augmented generation without requiring intermediate answer annotations. Achieve parameter-efficient performance matching much larger models through iterative preference optimization.

### Problem Context

Complex question answering demands multi-step reasoning and precise information retrieval. Traditional approaches struggle with:

- Decomposing ambiguous queries into retrievable subquestions
- Integrating multiple retrieved passages without losing coherence
- Requiring expensive intermediate supervision for each reasoning step
- Scaling linearly with model size for competitive performance

AceSearcher addresses these by training a single LLM to handle both question decomposition and answer synthesis through self-play, using only final-answer signals for optimization.

### Core Concept

AceSearcher treats reasoning as a cooperative process between two roles played by the same model:

**Decomposer (ρ)**: Transforms an input question into a sequence of subquestions designed to guide retrieval. Given question q, produces subquestions z = (z₁, z₂, ..., zₙ).

**Solver (π)**: Generates intermediate answers for each subquestion by consulting retrieved documents, then synthesizes all intermediate answers into a final response. Takes subquestion zᵢ, prior answers w<ᵢ, and retrieved context Dᵢ to produce intermediate answer wᵢ, then generates final answer a' from all components.

The key insight is that better decompositions enable the solver to retrieve more relevant context and produce higher-quality answers. This mutual dependency justifies joint optimization without separate roles or models.

### Architecture Overview

- **Unified Model**: Single LLM controlled by templates and prompts, no separate role-specific parameters
- **Two-Stage Training**: Supervised fine-tuning on diverse reasoning tasks, followed by preference-based reinforcement optimization
- **Preference Generation**: Multipath rollouts sample multiple decompositions and solution paths, creating preference pairs from max/min trajectories
- **Iterative DPO**: Direct Preference Optimization applied across three preference datasets (decomposition quality, intermediate answers, final answers) with convergence guarantees
- **Retrieval Integration**: Sparse retrieval allocates documents per subquestion (typically 15 total documents distributed across decomposed steps)

### Implementation

#### Stage 1: Supervised Fine-Tuning on Diverse Tasks

The foundation requires exposure to search, reasoning, and decomposition patterns across 180,600 examples:

```python
# sft_data_preparation.py
from dataclasses import dataclass
from typing import List

@dataclass
class SFTExample:
    instruction: str
    input: str
    output: str
    task_category: str  # "search", "reasoning", "decomposition"

# Data composition
sft_datasets = {
    "context_rich_qa": {
        "sources": ["NQ", "SQuAD", "DROP", "NarrativeQA", "Quoref", "ROPES", "FEVER", "TAT-QA"],
        "count": 132000,
        "purpose": "Extract answers from retrieved contexts"
    },
    "question_decomposition": {
        "sources": ["GSM8K", "ConvFinQA", "StrategyQA"],
        "count": 9600,
        "purpose": "Break complex questions into subquestions"
    },
    "chain_of_thought": {
        "sources": ["GSM8K", "TabMWP", "IfQA", "MathInstruct"],
        "count": 39000,
        "purpose": "Multi-step reasoning patterns"
    }
}

total_sft_examples = sum(v["count"] for v in sft_datasets.values())
# Result: 180,600 examples

# Hyperparameters for SFT
sft_config = {
    "batch_size": 64,
    "max_tokens": 2560,
    "epochs": 1,
    "warmup_steps_ratio": 0.05,
    "learning_rates": {
        "1.5B": 5e-6,
        "8B": 1e-6,
        "14B": 1e-6,
        "32B_LoRA": 1e-5
    },
    "optimizer": "AdamW",
    "lr_scheduler": "cosine"
}
```

Use standard next-token prediction loss. Mix tasks equally to ensure the model learns all three competencies.

#### Stage 2: Preference-Based Reinforcement Fine-Tuning

Generate preference pairs via multipath rollouts and optimize with iterative DPO:

```python
# preference_generation.py
import numpy as np
from typing import Tuple, List

class PreferenceDataGenerator:
    def __init__(self, model, retriever, num_decompositions=3, num_solutions=4):
        self.model = model
        self.retriever = retriever
        self.m = num_decompositions      # candidate decompositions per question
        self.m_prime = num_solutions      # candidate solutions per decomposition

    def compute_reward(self, question: str, answer: str, gold_answer: str) -> float:
        """
        Reward function: exact match with format validation
        r(q, a', a) = EM(a', a) × 𝕀(f(q, a') = 1)
        """
        em_score = 1.0 if self._normalize_answer(answer) == self._normalize_answer(gold_answer) else 0.0
        format_valid = self._check_format(question, answer)
        return em_score * float(format_valid)

From ADu2021/skillXiv, MIT.

More from ADu2021/skillXiv