AI Papers: Test-Time Scaling 핵심 논문 정리 — 추론 예산으로 성능을 끌어올리는 방법

from dataclasses import dataclass, field
from typing import Callable
import asyncio
import time

@dataclass
class ScoredCandidate:
    text: str
    score: float
    latency_ms: float
    token_count: int

@dataclass
class BestOfNConfig:
    n_samples: int = 4
    temperature: float = 0.7
    max_tokens: int = 1024
    timeout_seconds: float = 15.0
    cost_budget_usd: float = 0.05

class BestOfNPipeline:
    """프로덕션용 Best-of-N 파이프라인.

    후보 생성, 스코어링, 예산 제어를 하나의 흐름으로 관리한다.
    """

    def __init__(
        self,
        generator: Callable,
        scorers: list[Callable],
        config: BestOfNConfig = BestOfNConfig(),
    ):
        self.generator = generator
        self.scorers = scorers
        self.config = config

    async def generate_candidates(self, prompt: str) -> list[ScoredCandidate]:
        """N개의 후보를 비동기로 생성하고 스코어링한다."""
        candidates = []
        tasks = [
            self._generate_one(prompt)
            for _ in range(self.config.n_samples)
        ]

        results = await asyncio.gather(*tasks, return_exceptions=True)

        for result in results:
            if isinstance(result, Exception):
                continue  # timeout이나 API 에러는 건너뜀
            text, latency_ms, token_count = result
            score = self._compute_score(text)
            candidates.append(ScoredCandidate(
                text=text,
                score=score,
                latency_ms=latency_ms,
                token_count=token_count,
            ))

        return sorted(candidates, key=lambda c: c.score, reverse=True)

    def _compute_score(self, text: str) -> float:
        """여러 scorer의 가중 평균을 계산한다."""
        if not self.scorers:
            return 0.0
        scores = [scorer(text) for scorer in self.scorers]
        return sum(scores) / len(scores)

    async def _generate_one(self, prompt: str):
        """단일 후보 생성. timeout 초과 시 asyncio.TimeoutError 발생."""
        start = time.monotonic()
        result = await asyncio.wait_for(
            self.generator(prompt, temperature=self.config.temperature),
            timeout=self.config.timeout_seconds,
        )
        elapsed_ms = (time.monotonic() - start) * 1000
        return result.text, elapsed_ms, result.token_count

    async def run(self, prompt: str) -> ScoredCandidate | None:
        """최적 후보를 반환한다. 후보가 없으면 None."""
        candidates = await self.generate_candidates(prompt)
        return candidates[0] if candidates else None

from collections import Counter
import re
from typing import Optional


def extract_final_answer(text: str) -> Optional[str]:
    """CoT 출력에서 최종 답을 추출한다.

    여러 패턴을 시도하며, 하나도 매칭되지 않으면
    마지막 줄을 답으로 사용한다.
    """
    # 패턴 1: "정답: X", "Answer: X" 형태
    patterns = [
        r"(?:정답|답|answer|result)\s*[:：=]\s*(.+)",
        r"\\boxed\{(.+?)\}",              # LaTeX boxed 형태
        r"####\s*(.+)",                     # GSM8K 형태
        r"(?:따라서|그러므로|therefore)\s*(.+)",
    ]
    for pattern in patterns:
        matches = re.findall(pattern, text, re.IGNORECASE)
        if matches:
            return matches[-1].strip()

    # 모든 패턴 실패 시 마지막 비어있지 않은 줄 반환
    lines = [line.strip() for line in text.strip().splitlines() if line.strip()]
    return lines[-1] if lines else None


def normalize_answer(answer: str) -> str:
    """답 비교를 위한 정규화.

    공백, 구두점, 단위 표기 등의 사소한 차이를
    제거해서 동일한 답을 같은 것으로 인식한다.
    """
    s = answer.strip().lower()
    # 통화/단위 기호 제거
    s = re.sub(r"[$€₩\\,]", "", s)
    # 후행 마침표/쉼표 제거
    s = s.rstrip(".,;")
    # 다중 공백을 단일 공백으로
    s = re.sub(r"\s+", " ", s)
    return s


def self_consistency_vote(
    samples: list[str],
    min_agreement: int = 2,
) -> tuple[Optional[str], int, int]:
    """Self-Consistency 다수결을 수행한다.

    Returns:
        (best_answer, vote_count, total_valid):
        - best_answer: 가장 많이 등장한 답 (None if 추출 실패)
        - vote_count: 해당 답의 득표 수
        - total_valid: 답 추출에 성공한 전체 샘플 수
    """
    extracted = []
    for sample in samples:
        ans = extract_final_answer(sample)
        if ans is not None:
            extracted.append(normalize_answer(ans))

    if not extracted:
        return None, 0, 0

    counter = Counter(extracted)
    best_answer, vote_count = counter.most_common(1)[0]

    if vote_count < min_agreement:
        return None, vote_count, len(extracted)

    return best_answer, vote_count, len(extracted)

from dataclasses import dataclass, field
from typing import Optional
import heapq

@dataclass
class ThoughtNode:
    """추론 트리의 한 노드."""
    content: str
    depth: int
    score: float = 0.0
    parent: Optional['ThoughtNode'] = None
    children: list['ThoughtNode'] = field(default_factory=list)

    def path(self) -> list[str]:
        """루트부터 현재 노드까지의 추론 경로를 반환한다."""
        nodes = []
        current = self
        while current is not None:
            nodes.append(current.content)
            current = current.parent
        return list(reversed(nodes))


class BeamSearchReasoner:
    """Beam Search 기반 추론 탐색기.

    각 depth에서 beam_width만큼의 후보를 유지하며
    max_depth까지 탐색한 뒤 최고 점수 경로를 반환한다.
    """

    def __init__(
        self,
        expand_fn,       # (context: str) -> list[str]: 다음 추론 단계 후보들
        evaluate_fn,     # (partial_solution: str) -> float: 중간 평가 점수
        beam_width: int = 3,
        max_depth: int = 5,
    ):
        self.expand_fn = expand_fn
        self.evaluate_fn = evaluate_fn
        self.beam_width = beam_width
        self.max_depth = max_depth

    def search(self, problem: str) -> list[str]:
        """beam search로 최적 추론 경로를 찾는다."""
        root = ThoughtNode(content=problem, depth=0)
        beam = [root]

        for depth in range(1, self.max_depth + 1):
            all_candidates = []

            for node in beam:
                context = "\n".join(node.path())
                expansions = self.expand_fn(context)

                for thought in expansions:
                    child = ThoughtNode(
                        content=thought,
                        depth=depth,
                        parent=node,
                    )
                    partial = "\n".join(child.path())
                    child.score = self.evaluate_fn(partial)
                    node.children.append(child)
                    all_candidates.append(child)

            # 상위 beam_width개만 유지
            beam = heapq.nlargest(
                self.beam_width,
                all_candidates,
                key=lambda n: n.score,
            )

            if not beam:
                break

        # 전체 탐색 결과 중 최고 점수 경로 반환
        best = max(beam, key=lambda n: n.score)
        return best.path()

import subprocess
import tempfile
import textwrap
from pathlib import Path
from dataclasses import dataclass


@dataclass
class CodeEvalResult:
    code: str
    passed: bool
    tests_passed: int
    tests_total: int
    error_output: str
    execution_time_ms: float


def evaluate_code_candidate(
    code: str,
    test_code: str,
    timeout_seconds: int = 30,
) -> CodeEvalResult:
    """코드 후보를 테스트 스위트로 평가한다.

    실제 pytest를 실행하여 통과/실패를 판별하며,
    timeout, import error 등의 예외도 처리한다.
    """
    import time

    with tempfile.TemporaryDirectory() as tmpdir:
        solution_path = Path(tmpdir) / "solution.py"
        test_path = Path(tmpdir) / "test_solution.py"

        solution_path.write_text(code, encoding="utf-8")

        # test 파일에서 solution을 import할 수 있도록 구성
        full_test = textwrap.dedent(f"""\
            import sys
            sys.path.insert(0, "{tmpdir}")
            from solution import *

            {test_code}
        """)
        test_path.write_text(full_test, encoding="utf-8")

        start = time.monotonic()
        try:
            result = subprocess.run(
                ["python", "-m", "pytest", str(test_path), "-v", "--tb=short"],
                capture_output=True,
                text=True,
                timeout=timeout_seconds,
                cwd=tmpdir,
            )
            elapsed_ms = (time.monotonic() - start) * 1000

            # pytest 출력에서 passed/failed 수 파싱
            passed = result.returncode == 0
            lines = result.stdout.splitlines()
            summary_line = [l for l in lines if "passed" in l or "failed" in l]

            return CodeEvalResult(
                code=code,
                passed=passed,
                tests_passed=result.stdout.count("PASSED"),
                tests_total=result.stdout.count("PASSED") + result.stdout.count("FAILED"),
                error_output=result.stderr if not passed else "",
                execution_time_ms=elapsed_ms,
            )
        except subprocess.TimeoutExpired:
            elapsed_ms = (time.monotonic() - start) * 1000
            return CodeEvalResult(
                code=code,
                passed=False,
                tests_passed=0,
                tests_total=0,
                error_output=f"Execution timed out after {timeout_seconds}s",
                execution_time_ms=elapsed_ms,
            )


def pick_best_code(
    candidates: list[str],
    test_code: str,
) -> CodeEvalResult | None:
    """여러 코드 후보 중 가장 많은 테스트를 통과하는 것을 선택한다.

    전부 통과하는 후보가 여러 개면 실행 시간이 짧은 것을 우선한다.
    """
    results = [evaluate_code_candidate(c, test_code) for c in candidates]

    # 1순위: 통과 테스트 수, 2순위: 실행 시간 (낮을수록 좋음)
    results.sort(key=lambda r: (r.tests_passed, -r.execution_time_ms), reverse=True)

    return results[0] if results else None

# inference_policy.yaml
# 문제 난이도별 TTS 예산 할당 정책

policies:
  easy:
    description: '단순 QA, FAQ, 분류'
    strategy: single_pass
    samples: 1
    temperature: 0.0
    max_tokens: 256
    estimated_cost_per_query_usd: 0.001

  medium:
    description: '요약, 번역, 간단한 추론'
    strategy: best_of_n
    samples: 3
    temperature: 0.5
    max_tokens: 512
    scorer: format_and_length
    estimated_cost_per_query_usd: 0.005

  hard:
    description: '수학, 코딩, 멀티스텝 추론'
    strategy: self_consistency
    samples: 8
    temperature: 0.7
    max_tokens: 1024
    min_agreement: 3
    estimated_cost_per_query_usd: 0.015

  critical:
    description: '코드 리뷰, 안전 판단, 의료/법률'
    strategy: verified_best_of_n
    samples: 5
    temperature: 0.6
    max_tokens: 2048
    verifier: domain_specific_verifier
    fallback_on_low_confidence: human_review
    estimated_cost_per_query_usd: 0.025

routing:
  classifier: difficulty_classifier_v2
  fallback_policy: medium

guardrails:
  max_total_tokens_per_query: 16384
  timeout_seconds: 30
  daily_budget_usd: 500
  alert_threshold_usd: 400

-- TTS 운영 대시보드용 쿼리
-- 일별 전략별 성능/비용/지연을 한 눈에 파악한다

SELECT
    date_trunc('day', created_at)     AS day,
    tts_strategy,
    COUNT(*)                          AS total_queries,
    AVG(correct::int)                 AS accuracy,
    PERCENTILE_CONT(0.50) WITHIN GROUP (ORDER BY latency_ms) AS p50_latency_ms,
    PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY latency_ms) AS p95_latency_ms,
    AVG(samples_generated)            AS avg_samples,
    SUM(cost_usd)                     AS total_cost_usd,
    AVG(cost_usd)                     AS avg_cost_per_query,
    -- 효율성 지표: 정답 1건당 비용
    SUM(cost_usd) / NULLIF(SUM(correct::int), 0) AS cost_per_correct_answer
FROM reasoning_eval_runs
WHERE created_at >= CURRENT_DATE - INTERVAL '14 days'
GROUP BY 1, 2
ORDER BY 1 DESC, 2;