AetherSearch_DPO

<div align="center"> <img src="assets/aethersearch-mark.svg" alt="AetherSearch monogram" width="144">

🔭 AetherSearch DPO

A preference-aligned search agent for retrieval-augmented reasoning

Built on AetherSearch SFT and aligned with Direct Preference Optimization (DPO) over 2,126 search-trajectory preference pairs.

<p> <img src="https://img.shields.io/badge/Method-DPO-0F766E?style=flat-square" alt="Training method: DPO"> <img src="https://img.shields.io/badge/Pairs-2%2C126-F59E0B?style=flat-square" alt="Training pairs: 2,126"> <img src="https://img.shields.io/badge/Context-32K-2563EB?style=flat-square" alt="Context window: 32K"> </p>

🏠 Project · 🧠 SFT checkpoint

</div>

🔌 Bring your own retriever. AetherSearch DPO is a search-agent policy, not a self-contained question-answering service. The host runtime must execute each <search>...</search> request and return evidence inside <information>...</information>.

✨ Highlights

• 🎯 Preference aligned — favors stronger search decisions, evidence use, and final answers over paired alternatives.
• 🔎 Search native — produces explicit reasoning, retrieval requests, and grounded answers using the AetherSearch XML protocol.
• 🔁 Multi-step retrieval — can request additional evidence when an existing trajectory is insufficient.
• 🧾 Retrieval-aware DPO — retrieved information remains visible as context but is excluded from preference-loss scoring.

🧠 How it works

Question
   │
   ▼
<think>reason about what is missing</think>
   │
   ▼
<search>focused retrieval query</search> ─────► Search / RAG backend
   ▲                                                  │
   └──── <information>retrieved evidence</information> ◄────┘
   │
   ├── repeat the search loop when more evidence is needed
   ▼
<answer>evidence-grounded final answer</answer>

The model produces reasoning, search, and answer spans. The surrounding runtime parses each completed <search> span, runs retrieval, appends the result as <information>, and resumes generation until the model emits an <answer> span.

📊 Model at a glance

🎯 DPO alignment

Supervised fine-tuning teaches the model how to follow the search protocol; DPO then teaches it which of two valid-looking continuations is preferable. For a shared prompt \(x\), preferred continuation \(y_w\), rejected continuation \(y_l\), policy \(\pi_\theta\), and frozen SFT reference \(\pi_{\mathrm{ref}}\), training minimizes:

$$ \mathcal{L}{\mathrm{DPO}} = -\log \sigma!\left( \beta \left[ \log \frac{\pi\theta(y_w \mid x)}{\pi_{\mathrm{ref}}(y_w \mid x)}

\log \frac{\pi_\theta(y_l \mid x)}{\pi_{\mathrm{ref}}(y_l \mid x)} \right] \right). $$

In plain terms, the policy learns to widen the preference margin between the chosen and rejected search trajectories while the frozen SFT model anchors the update. This directly optimizes pairwise preferences without training a separate reward model or running an online RL loop.

The loss is adapted to the agent-environment boundary:

• Shared prompt tokens are masked on both sides.
• Environment-provided <information>...</information> spans inside either continuation are masked while remaining visible as context.
• All remaining assistant tokens contribute to summed sequence log probabilities.
• Answer-terminal continuations include the final <|im_end|> token; search-terminal continuations stop before it so the runtime can insert the next retrieval result.

Both the initial policy and frozen reference use the pinned AetherSearch SFT checkpoint. The released run uses \(\beta=0.1\).

🛠️ Training recipe

The policy and frozen reference are both sharded with ZeRO-3; gradients and optimizer updates are applied only to the policy.

🚀 Quick start

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "muradil211/AetherSearch_DPO"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
model.config.use_cache = True

Loading the checkpoint is only the first step. For end-to-end use, stop generation after each complete <search> request, execute it with your retriever, append the result as <information>, and resume generation. Stop when the model emits a complete <answer> span, and preserve the XML protocol exactly throughout the loop.

⚠️ Limitations

• DPO improves fit to the supplied preference pairs; it does not guarantee factual, supported, or safe outputs.
• Preference labels include curated hard negatives and trajectory corrections, not human preference votes for every pair.
• Retrieval quality, evidence validation, answer verification, and deployment safeguards remain the caller's responsibility.

📜 Terms

No additional blanket license is asserted here. Review the Qwen2.5-3B-Instruct license, the AetherSearch SFT terms, and the AetherSearch DPO data attribution before redistribution or downstream use.

---

<div align="center">

Built for agentic search and retrieval-augmented reasoning. 🔎✨

</div>