튜토리얼: AIME(수학)를 위한 GEPA

튜토리얼: AIME(수학)를 위한 GEPA (GEPA for AIME)

이 튜토리얼에서는 GPT-4.1 Mini의 Chain of Thought(dspy.ChainOfThought)를 dspy.GEPA 옵티마이저로 최적화해 수학 문제(AIME)를 풀어볼게요!

출처: 문서

본문

AIME 데이터셋 로드 (Loading the AIME dataset)

AIME 시험은 각 연도마다 15문제짜리 문제 세트 2개로 구성돼요. 이 튜토리얼에서는 최적화를 위해 이전 연도(2022-2024)의 AIME 문제 세트를 사용할 거예요(총 3년 × 2세트 × 15문제 = 90문제, 학습/검증 세트에 균등 분할). 그리고 AIME 2025(2세트 × 15문제 = 30문제)에서 성능을 테스트합니다. AIME 2025는 작은 세트이므로 평가의 통계적 안정성을 위해 5번 반복해요.

import dspy
from datasets import load_dataset

def init_dataset():
    train_split = load_dataset("AI-MO/aimo-validation-aime")['train']
    train_split = [
        dspy.Example({
            "problem": x['problem'],
            'solution': x['solution'],
            'answer': x['answer'],
        }).with_inputs("problem")
        for x in train_split
    ]
    import random
    random.Random(0).shuffle(train_split)
    tot_num = len(train_split)

    test_split = load_dataset("MathArena/aime_2025")['train']
    test_split = [
        dspy.Example({
            "problem": x['problem'],
            'answer': x['answer'],
        }).with_inputs("problem")
        for x in test_split
    ]

    train_set = train_split[:int(0.5 * tot_num)]
    val_set = train_split[int(0.5 * tot_num):]
    test_set = test_split * 5

    return train_set, val_set, test_set

train_set, val_set, test_set = init_dataset()

len(train_set), len(val_set), len(test_set)

(45, 45, 150)

프로그램 정의: 간단한 `dspy.ChainOfThought`

class GenerateResponse(dspy.Signature):
    """Solve the problem and provide the answer in the correct format."""
    problem = dspy.InputField()
    answer = dspy.OutputField()

program = dspy.ChainOfThought(GenerateResponse)

최적화되지 않은 Chain Of Thought 평가하기 (Evaluating unoptimized Chain Of Thought)

import dspy
evaluate = dspy.Evaluate(
    devset=test_set,
    metric=metric,
    num_threads=32,
    display_table=True,
    display_progress=True
)

evaluate(program)

Average Metric: 70.00 / 150 (46.7%): 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 150/150 [00:01<00:00, 119.75it/s]
2025/08/12 21:49:36 INFO dspy.evaluate.evaluate: Average Metric: 70 / 150 (46.7%)

EvaluationResult(score=46.67, results=<list of 150 results>)

최적화되지 않은 CoT 프로그램은 46.67%의 정확도를 보여줍니다.

GEPA로 최적화 (Optimize with GEPA)

dspy.GEPA 최적화기를 사용하고, GEPA가 피드백으로 사용할 수 있는 텍스트 피드백을 제공하는 지표를 설정합니다. GEPA는 프로그램의 궤적을 반성하고 새 프롬프트를 제안해요. 최적화 후 평가에 사용되는 것과 동일한 metric을 최적화 지표로 사용할 수 있으며, 반성용 LM(reflection_lm)을 구성합니다.

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

GEPA로 최적화된 Chain Of Thought 평가하기 (Evaluating the Chain Of Thought optimized with GEPA)

evaluate(optimized_program)

Average Metric: 85.00 / 150 (56.7%): 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 150/150 [00:00<00:00, 476.89it/s]
2025/08/12 23:53:14 INFO dspy.evaluate.evaluate: Average Metric: 85 / 150 (56.7%)

EvaluationResult(score=56.67, results=<list of 150 results>)

최적화 후 46.67%에서 56.67%로 약 10퍼센트포인트 향상됐어요. 즉 GEPA는 GPT-4.1 Mini로 AIME 2025에서 큰 성과를 냈습니다.

생성된 프롬프트 보기 (Let's see the prompt generated)

print(optimized_program.predict.signature.instructions)

GEPA가 생성한 최적화된 지침은 문제 유형을 파싱하고, 도메인 제약을 적용하며, 대수적 항등식과 모듈러 수학을 사용해 탐색 공간을 줄이는 등 상세한 문제 해결 전략을 담고 있어요. 예를 들어:

You will be given one math problem as plain text under a key like "problem." Your job is to solve it correctly and return:

- reasoning: a concise, logically ordered solution that uses identities/structure to avoid brute force, ends with a quick verification.
- answer: the final requested number/expression only (no extra words).

Formatting:
- Use exactly two top-level fields named "reasoning" and "answer."
- Keep reasoning succinct but complete. Bullet points are fine.
- The answer field must contain only the final value requested (e.g., 227, 585, 601).

General problem-solving guidance:
- Parse the problem type (e.g., base representation, intersecting families of subsets, avoiding arithmetic progressions, symmetric sums with constraints, ordered tuples counting).
- Always enforce domain constraints (e.g., base-b digits in 0..b−1; no leading zero for base-10 "three-digit"; ordered vs unordered families; strict increase conditions in sequences).
- Use algebraic identities and modular arithmetic to reduce the search space; prefer structural arguments over naive enumeration.
- For "greatest/least" questions, derive tight bounds and give a construction that attains them.
...

이 지침은 GEPA가 수학 문제 해결의 미묘한 함정(기수 변환, 십진수 회문, 고정된 합의 대칭 합, 교차 부분집합족, 산술 진행 회피 등)을 세밀하게 포착해 프롬프트에 반영했음을 보여줘요.