muradil211/AetherSearch_DPO

🔭 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.

🏠 Project · 🧠 SFT checkpoint

🔌 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

Field	Value
🧬 Base checkpoint	muradil211/AetherSearch_SFT
🧬 Base revision	437aca474d3966e57e82af565db95d0ad64aa24d
🏗️ Architecture	Qwen2 causal language model
🔢 Parameters	3,085,938,688
🎛️ Weight dtype	BF16
📏 Context window	32,768 positions; training sequences capped at 4,096
🎯 Alignment method	Direct Preference Optimization

🎯 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 $xx$, preferred continuation $y_{w}y\_w$, rejected continuation $y_{l}y\_l$, policy ${\pi }_{\theta }\backslash pi\_\backslash theta$, and frozen SFT reference ${\pi }_{\mathrm{r}\mathrm{e}\mathrm{f}}\backslash pi\_\{\backslash mathrm\{ref\}\}$, training minimizes:

$${\mathcal{L}}_{\mathrm{D}\mathrm{P}\mathrm{O}}=-\log \sigma \text{\hspace{0.17em}⁣}\left(\beta [\log \frac{{\pi }_{\theta }(y_w\mid x)}{{\pi }_{\mathrm{r}\mathrm{e}\mathrm{f}}(y_w\mid x)}-\log \frac{{\pi }_{\theta }(y_l\mid x)}{{\pi }_{\mathrm{r}\mathrm{e}\mathrm{f}}(y_l\mid x)}]\right)\mathrm{.}\backslash mathcal\{L\}\_\{\backslash mathrm\{DPO\}\}\; =\; -\backslash log\; \backslash sigma\backslash !\backslash left(\; \backslash beta\; \backslash left[\; \backslash log\; \backslash frac\{\backslash pi\_\backslash theta(y\_w\; \backslash mid\; x)\}\{\backslash pi\_\{\backslash mathrm\{ref\}\}(y\_w\; \backslash mid\; x)\}\; -\; \backslash log\; \backslash frac\{\backslash pi\_\backslash theta(y\_l\; \backslash mid\; x)\}\{\backslash pi\_\{\backslash mathrm\{ref\}\}(y\_l\; \backslash mid\; x)\}\; \backslash right]\; \backslash 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\backslash beta=0.1$.

🛠️ Training recipe

Setting	Value	Setting	Value
Epochs	1	Learning rate	5e-7
DPO beta	0.1	Scheduler	Cosine
Global batch size	12 pairs	Per-device batch	1 pair
Precision	BF16	Max sequence length	4,096
Warmup ratio	0.03	Weight decay	0.0
Distributed optimizer	DeepSpeed ZeRO-3	Seed	42

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.

---

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