튜토리얼: 다중 홉 연구를 위한 온라인 RL

튜토리얼: 다중 홉 연구를 위한 온라인 RL (Online RL for Multi-Hop Research)

경고: 이 기능은 새롭고 매우 실험적입니다. DSPy의 거의 모든 다른 것과 달리, 현재는 순수한 개념 증명·개발 모드에 있으며, 커뮤니티 참여를 장려하기 위해 배포합니다.

이 튜토리얼에는 DSPy의 Arbor RL 프레임워크도 필요하며, 아래와 같이 설치할 수 있어요:

> pip install -U arbor-ai

또한 DSPy를 main 브랜치에서 설치해야 할 수도 있어요:

> pip install -U git+https://github.com/stanfordnlp/dspy.git@main

출처: 문서

본문

import dspy
import arbor
from arbor import ArborGRPO, ArborProvider
arbor_server_info = arbor.init() # Initialize the Arbor server in the background

port = 7453
local_lm_name = "Qwen/Qwen2.5-1.5B-Instruct"
local_lm = dspy.LM(
    model=f"openai/arbor:{local_lm_name}",
    provider=ArborProvider(),
    api_base=arbor_server_info["base_url"],
    # Arbor checks to make sure these match the training config
    temperature=1.0,
    top_p=1.0,
    top_k=-1,
    repetition_penalty=1.0,
    max_tokens=2048,
)

dspy.configure(lm=local_lm)

의존성 설치 및 데이터 다운로드 (Install dependencies and download data)

검색을 위해 멋진 BM25S 라이브러리를 사용할 거예요. 꽤 가볍기 때문이에요. 이 구성 요소는 원하는 무엇으로든 바꿀 수 있어요.

> pip install -U bm25s PyStemmer "jax[cpu]"

다음으로 2017년 기준 모든 5,000,000개의 Wikipedia 페이지의 스냅샷 초록(즉, 첫 문단)을 다운로드할 거예요. 이를 검색 코퍼스로 사용합니다.

이는 압축 상태 500MB이므로, 다운로드와 압축 해제에 2-3분이 걸릴 수 있어요.

from dspy.utils import download

download("https://huggingface.co/dspy/cache/resolve/main/wiki.abstracts.2017.tar.gz")
!tar -xzvf wiki.abstracts.2017.tar.gz

그다음 BM25 검색을 위해 인덱싱해 볼게요! 2-3분이 걸립니다.

import orjson
import bm25s
import Stemmer

corpus = []

with open("wiki.abstracts.2017.jsonl") as f:
    for line in f:
        line = orjson.loads(line)
        corpus.append(f"{line['title']} | {' '.join(line['text'])}")

stemmer = Stemmer.Stemmer("english")
corpus_tokens = bm25s.tokenize(corpus, stopwords="en", stemmer=stemmer)

retriever = bm25s.BM25(k1=0.9, b=0.4)
retriever.index(corpus_tokens)

HoVer 데이터셋 로드 (Load the HoVer dataset)

작업을 위한 데이터셋을 로드해 볼게요. HoVer 다중 홉 작업의 예를 로드하는데, 여기서 입력은 (정말로!) 복잡한 주장이고 우리가 찾는 출력은 해당 주장을 사실 확인(fact-check)하는 데 필요한 Wikipedia 페이지 집합입니다.

데이터셋의 이전 버전을 설치해야 제대로 동작할 수도 있어요...

> pip install datasets==3.6.0

import random
from dspy.datasets import DataLoader

kwargs = dict(fields=("claim", "supporting_facts", "hpqa_id", "num_hops"), input_keys=("claim",))
hover = DataLoader().from_huggingface(dataset_name="hover-nlp/hover", split="train", trust_remote_code=True, **kwargs)

hpqa_ids = set()
hover = [
    dspy.Example(claim=x.claim, titles=list(set([y["key"] for y in x.supporting_facts]))).with_inputs("claim")
    for x in hover
    if x["num_hops"] == 3 and x["hpqa_id"] not in hpqa_ids and not hpqa_ids.add(x["hpqa_id"])
]

random.Random(0).shuffle(hover)
trainset, devset, testset = hover[:600], hover[600:900], hover[900:]
len(trainset), len(devset), len(testset)

이제 Wikipedia에서 검색하는 함수를 정의해 볼게요. 우리의 BM25 인덱스를 사용할 거예요.

def search(query: str, k: int) -> list[str]:
    tokens = bm25s.tokenize(query, stopwords="en", stemmer=stemmer, show_progress=False)
    results, scores = retriever.retrieve(tokens, k=k, n_threads=1, show_progress=False)
    run = {corpus[doc]: float(score) for doc, score in zip(results[0], scores[0])}
    return list(run.keys())

다중 홉 연구를 위한 DSPy 프로그램 (A DSPy program for multi-hop research)

이제 DSPy에서 다중 홉 프로그램을 정의해 볼게요. generate_query와 append_notes 모듈로 구성된 아주 간단한 프로그램이에요. 지침은 대체로 필요하지 않지만, 신중하게 정의할 거예요.

instr1 = """
Given a claim and some key facts, generate a follow-up search query to find the next most essential clue towards verifying or refuting the claim. The goal ultimately is to find all documents implicated by the claim.
""".strip()

instr2 = """
Given a claim, some key facts, and new search results, identify any new learnings from the new search results, which will extend the key facts known so far about the whether the claim is true or false. The goal is to ultimately collect all facts that would help us find all documents implicated by the claim.
"""


class ResearchHop(dspy.Module):
    def __init__(self, num_docs, num_hops):
        self.num_docs, self.num_hops = num_docs, num_hops
        self.generate_query = dspy.ChainOfThought(dspy.Signature("claim, key_facts -> followup_search_query", instr1))
        self.append_notes = dspy.ChainOfThought(dspy.Signature("claim, key_facts, new_search_results -> new_key_facts", instr2))

    def forward(self, claim: str) -> list[str]:
        key_facts = []
        retrieved_docs = []

        for hop_idx in range(self.num_hops):
            query = self.generate_query(claim=claim, key_facts=key_facts).followup_search_query if hop_idx else claim
            search_results = search(query, k=self.num_docs)
            retrieved_docs.extend(search_results)

            if hop_idx == self.num_hops - 1:
                break
                
            prediction = self.append_notes(claim=claim, key_facts=key_facts, new_search_results=search_results)
            key_facts.append(prediction.new_key_facts)

        return dspy.Prediction(key_facts=key_facts, retrieved_docs=retrieved_docs)

이 작업에서 성공을 정의하는 지표 (Define metrics for success in this task)

def recall(example, pred, trace=None):
    gold_titles = example.titles
    retrieved_titles = [doc.split(" | ")[0] for doc in pred.retrieved_docs]
    return sum(x in retrieved_titles for x in set(gold_titles)) / len(gold_titles)

evaluate = dspy.Evaluate(devset=devset, metric=recall, num_threads=16, display_progress=True, display_table=5)

dspy.GRPO로 ResearchHop 시스템 최적화하기 (Optimize the ResearchHop system with dspy.GRPO)

program = ResearchHop(num_docs=4, num_hops=2)
program.set_lm(local_lm)

# NOTE: Training on 4 GPUs.
train_kwargs = {
    "per_device_train_batch_size": 2,
    "gradient_accumulation_steps": 24/6,
    "temperature": 1.0,
    "top_k": -1,
    "top_p": 1.0,
    "repetition_penalty": 1.0,
    "beta": 0.00,
    "learning_rate": 1e-6,
    "gradient_checkpointing": True,
    "bf16": True,
    "lr_scheduler_type": "constant_with_warmup",
    "loss_type": "dapo",
    "max_steps": 1000,
    "report_to": "wandb",
    "log_completions": True,
    "logging_steps": 1,
    "max_prompt_length": None,
    "max_completion_length": None,
    "scale_rewards": False,
    "max_grad_norm": 1.0,
    "lora_config": {
        "lora_alpha": 16,
        "lora_dropout": 0.05,
        "r": 8,
        "target_modules": ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"],
    },
    "num_training_gpus": 3,
    "num_inference_gpus": 1,
    "weight_decay": 0.001,
}

compiler = ArborGRPO(
    metric=recall,
    num_dspy_examples_per_grpo_step=6,
    num_rollouts_per_grpo_step=24,
    exclude_demos=True,
    num_train_steps=1000,
    num_threads=16,
    use_train_as_val=False,
    num_steps_for_val=50,
    train_kwargs=train_kwargs,
    checkpoint="single-best",
)

optimized_program = compiler.compile(
    student=program,
    trainset=trainset,
    valset=devset,
)

이제 GRPO가 적용된 프로그램을 사용할 수 있어요.

example = devset[0]
optimized_program(**example.inputs())

우리의 예비 실험에서, 약 18시간 훈련하면 (devset) recall이 61.8%에서 66.2%로 올라가요. 이는 비용/품질 기준으로 보면 dspy.MIPROv2나 dspy.SIMBA 같은 프롬프트 옵티마이저를 실행할 때보다 대체로 나쁘지만, 작은 LM을 위한 임의의 LM 프로그램에 대한 온라인 RL로는 여전히 매우 견고한 출발점입니다.

더 알아보기 (Learn more)

• Multi-Hop Search 튜토리얼
• RL PAPILLON 튜토리얼
• RL AI Program