기술 요약: Best-of-Better-N (BoBN)

문제 정의

Best-of-N (BoN)과 같은 추론 시 정렬(Inference-time alignment, ITA) 방법은 모델의 가중치를 수정하지 않고도 대규모 언어 모델(LLM)을 인간의 선호도나 특정 작업 요구사항에 정렬하는 것을 목표로 합니다. 이러한 방법들은 "샘플링 및 평가(sample-and-evaluate)" 프레임워크를 기반으로 작동합니다. 즉, 참조 정책 $\pi_{ref}$로부터 $N$개의 응답을 샘플링하고, 보상 모델 $r$에 따라 가장 높은 보상을 가진 응답을 선택합니다.

ITA의 효용성은 근본적으로 두 가지 가정에 의해 제약됩니다:

1. 커버리지 (A1): 참조 정책 $\pi_{ref}$가 적어도 하나의 고품질(고보상) 응답을 샘플링할 수 있어야 합니다.
2. 선택 (A2): 선택 기준이 샘플들 사이에서 고품질 응답을 성공적으로 식별해야 합니다.

현재의 ITA 연구는 주로 보상 해킹(reward hacking)을 완화하기 위해 선택 기준을 개선함으로써 (A2)를 향상시키는 데 집중하고 있습니다. 그러나 (A1)에는 여전히 결정적인 병목 현상이 남아 있습니다. 만약 $\pi_{ref}$가 다운스트림 작업에 정렬되어 있지 않다면(예: 유해한 프롬프트에 대해 거절하도록 안전 학습이 되지 않은 모델), 고보상 응답에 할당되는 확률이 극히 낮아집니다. 이 경우, 어떤 선택 전략도 성공할 수 없는데, 그 이유는 고보상 영역의 출력 공간이 참조 모델에 의해 사실상 접근 불가능하기 때문입니다. 기존의 솔루션들은 흔 대개 미세 조정(비용이 많이 들고 가중치 접근 권한이 필요함)을 요구하거나, 작업 특이적 관련성이 부족할 수 있는 정적인 인컨텍스트 예시(static in-context examples)를 사용합니다.

방법론: Best-of-Better-N (BoBN)

저자들은 가중치 접근이나 미세 조정 없이 $\pi_{ref}$의 샘플링 분포를 고보상 영역으로 이동시키기 위해 설계된 추론 프레임십인 **Best-of-Better-N (BoBN)**을 제안합니다. BoBN은 인컨텍스트 학습(In-Context Learning, ICL)을 검색(retrieval) 및 리스타일링(restyling) 파이프라인과 결합하여 활용합니다.

BoBN 워크플로우는 세 단계로 구성됩니다:

1. 검색 (Retrieval): 입력 프롬프트 $x$가 주어지면, 시스템은 임베딩 모델을 사용하여 데이터셋 $D$로부터 가장 유사한 프롬프트-응답 쌍 $(x_i, y_i)$ 상위 $K$개를 검색합니다. 이 예시들은 특정 쿼리에 관련된 고보상 행동의 데모 역할을 합니다.
2. 리스타일링 (Restyling): 검색된 예시들이 타겟 작업의 특정 제약, 형식, 페르소나와 일치하도록 하기 위해, 참조 모델 $\pi_{ref}$를 사용하여 검색된 응답 $y_i$를 "리스타일링"합니다. 구체적으로, $\pi_{ref}$는 " $y_i$를 스타일 $T$에 맞게 다시 작성하라"와 같은 프롬프트에 조건화되어 $y_i$를 $\hat{y}_i$로 재작성합니다. 이 단계는 검색된 데이터의 정렬 로직을 보존하면서, 컨텍스트와 참조 모델 간의 토큰 수준 분포 일관성을 보장합니다.
3. 생성 및 선택 (Generation and Selection): 리스타일링된 예시들 $C_K = \{(x_i, \hat{y}_i)\}$를 인컨텍스트 예시로서 입력 프롬프트 $x$ 앞에 붙입니다. 그런 다음 참조 모델은 조건부 분포 $\pi_{ref}(\cdot | C_K, x)$로부터 $N$개의 응답을 샘플링합니다. 마지막으로, 선택 기준(예: BoN)이 보상 모델을 기반으로 최적의 응답을 선택합니다.

$C_K$에 조건화됨으로써, BoBN은 $\pi_{ref}$의 확률 질량을 고보상 영역으로 효과적으로 이동시켜, 고품질 후보가 처음부터 생성될 가능성을 높입니다.

주요 기여

1. 새로운 프레임워크: 검색 및 리스타일링된 예시를 활용하여 고품질 응답에 대한 $\pi_{ref}$의 커버리지를 개선하는 ICL 기반 접근 방식인 BoBN을 제 제안하였습니다.
2. 이론적 분석: 검색 품질이 커버리지와 다운스트림 성능을 어떻게 제어하는지에 대한 이론적 특성을 규명했습니다. 저자들은 작업 관련 검색이 정적 ICL이나 ICL 미사용 시보다 커버리지(골든 정책과 참조 정책 사이의 카이제곱 발산으로 정의됨)를 입증 가능하게 감소시킨다는 것을 보여줍니다. 이러한 커버리지 감소가 BoN 선택과 결합되어 작업 성능을 향상시킨다는 점을 증명했습니다.
3. 실증적 검증: 안전 정렬(StrongReject, SORRY-Bench) 및 수학적 추론(GSM8k, MATH500) 벤치마크에 대한 광범위한 실험을 수행했습니다. 결과에 따르면 BoBN은 여러 참조 LLM(zephyr-7b, tulu3-8b)에서 $N=1$일 때(개별 응답 품질)와 $N>1$일 때(선택과 결합 시) 모두 성능을 향상시켰습니다.

실험 결과

저자들은 표준 BoN($\pi_{ref}$를 직접 사용) 및 URIAL(정적 예시를 사용하는 베닐라 ICL)과 BoBN을 비교 평가했습니다.

• 안전 정렬 (Safety Alignment): StrongReject 벤치마크에서 BoBN은 베이스라인들을 크게 앞질렀습니다. zephyr-7b 모델의 경우, $N=16$인 BoBN은 0.033의 점수(최적에 근접)를 달取得하여 Claude Sonnet-3.7과 같은 프런티어 모델과 대등한 성능을 보인 반면, $\pi_{ref}$ 단독으로는 0.425를 기록했습니다. BoBN은 또한 URIAL 및 $\pi_{ref}$에 비해 더 우수한 커버리지(더 낮은 $\log C$)를 보여주었습니다.
• 수학적 추론 (Mathematical Reasoning): GSM8k 및 MATH500에서 BoBN은 일관된 정확도 향상을 가져왔습니다. GSM8k에 대한 zephyr-7b의 경우, BoBN은 $N=16$에서 58.7%의 정확도를 달성하여, BoN을 적용한 $\pi_{ref}$의 43.9%보다 유의미한 이득을 보였습니다. 특히, $N=4$인 BoBN이 종종 $N=16$인 베이스라인들의 성능과 일치하거나 이를 능가했는데, 이는 높은 샘플링 효율성을 나타냅니다.
• 비용 및 지연 시간 (Cost and Latency): BoBN은 생성 전 $\pi_{ref}$ 호출 $K$번을 요구하는 리스타일링 단계로 인해 추가적인 지연 시간이 발생하지만, 상당한 토큰 효율성을 제공합니다. $N=4$인 BoBN은 유사하거나 더 나은 성능을 달성하면서도 $N=16$인 URIAL에 비해 토큰 사용량을 최대 50%까지 절감했습니다.
• 절제 연구 (Ablation Studies):
  • 리스타일링 (Restyling): 리스타일링 단계를 제거하면 안전 정렬 성능이 약화되고 커버리지가 높아졌으며(성능 악화), 이는 검색된 예시를 타겟 작업의 스타일에 맞추는 과정의 필요성을 확인시켜 줍니다.
  • 검색 개수 ($K$): $K$가 1에서 8로 증가함에 따라 성능이 향 향상되었으나, $K=16$이 되면 오히려 저하되었습니다. 이는 너무 많은 예시를 검색하면 덜 관련 있는 데이터가 유입되어 인컨텍스트 신호를 희석시킨다는 것을 시사합니다.

의의 및 주장

본 논문은 BoBN이 추론 시 정렬의 근본적인 한계, 즉 정렬되지 않은 참조 모델이 고보상 응답을 샘플링하지 못하는 문제를 해결한다고 주장합니다. 검색 및 리스타일링된 인컨텍스트 예시를 통해 샘플링 분포를 이동시킴으로써 BoBN은 다음을 가능하게 합니다:

• 고정된 예산 내 높은 품질: 고정된 샘플 수 $N$에 대해 더 나은 작업 성능을 제공합니다.
• 효율성: 목표 성능을 달성하기 위해 더 작은 $N$을 사용하여 추론 비용을 줄입니다.
• 훈련이 필요 없는 정렬 (Training-Free Alignment): 다운스트림 작업을 위해 명시적으로 미세 조정되지 않은 모델의 정렬을, 모델의 가중치나 라이클리후드(likelihood)에 대한 접근 없이도 개선합니다.

저자들은 선형 트랜스포머(linear transformers)에 기반한 이론적 분석을 통해, 검색 품질, 커버리지 감소, 그리고 다운스트림 작업 성능 사이의 입증 가능한 연결 고리를 제공하며, 단순한 선택 휴리스틱을 넘어 ITA를 개선하기 위한 원칙적인 접근 방식을 제시한다고 강조합니다.