Best-of-N 은 최선인가: 추론 시점 정렬의 이론적 한계와 대안 알고리즘 Inference-Time Pessimism 논문 리뷰

핵심 메시지

• 최근 LLM 연구는 추가 학습 대신 추론 시점에 연산을 더 써서 답변 품질을 올리는 방향, 즉 추론 시점 정렬(inference-time alignment)에 집중하고 있다.
• 가장 단순하고 널리 쓰이는 Best-of-N 은 이 논문에서 처음으로 이론적으로 정밀하게 분석됐고, 도달할 수 있는 최선(스카이라인)에 미치지 못한다는 사실이 증명됐다.
• 샘플 수 N 을 키울수록 성능이 오히려 나빠지는 현상은 그동안 실험적으로만 보고됐는데, 이 논문이 리워드 모델 과최적화 항으로 이론적 설명을 제시했다.
• 저자들은 카이제곱 정규화 강화학습 목적함수의 최적 정책에서 표본을 뽑는 Inference-Time Pessimism 알고리즘을 제안해, 스카이라인과 같은 차수의 보장을 얻었다.
• 실험에서도 Best-of-N 은 N 이 커지면 성능이 꺾이지만 제안 알고리즘은 일정 수준 이상에서 성능이 유지되거나 개선됐다.

쉽게 이해하기

논문 읽기 모임 PR-12 의 519번째 발표로 소개된 이 논문의 주제는 추론 시점 정렬이다. 이미 학습을 마친 모델을 더 훈련시키는 대신, 추론할 때 연산을 더 들여 답변 품질을 끌어올리는 접근으로 테스트 타임 컴퓨트라고도 불린다. 발표자는 마지막 저자가 강화학습 이론 분야에서 활발히 논문을 내는 마이크로소프트 리서치 소속 연구자라고 소개하며, 이론과 실용을 함께 다루는 논문이라고 설명한다.

설정은 단순하다. 우리에게는 이미 학습된 정책, 즉 언어 모델이 있고, 선호 데이터로 학습한 리워드 모델이 있으며, 실제로는 알 수 없는 참된 리워드가 따로 존재한다고 가정한다. 프롬프트가 들어오면 모델에서 응답을 여러 개 뽑고, 각 응답의 리워드 값과 확률 정보를 조합해 최종 응답 하나를 내놓는다. 여기서 두 가지 질문이 생긴다. 리워드 모델이 완벽하지 않은데 우리는 얼마나 최적에 가까이 갈 수 있는가, 그리고 그러기 위해 샘플을 몇 개나 뽑아야 하는가.

논문은 먼저 하한을 제시한다. 이상적인 답의 성능과 실제 출력의 성능 차이는 아무리 잘해도 어떤 값 아래로 내려갈 수 없는데, 그 값은 참조 정책이 목표 정책을 얼마나 잘 덮는지를 나타내는 커버리지와 리워드 모델의 부정확성으로 결정된다. 저자들은 이 도달 가능한 최선을 비유적으로 스카이라인이라고 부른다. 그런 다음 Best-of-N 의 후회(regret) 상한을 구하는데, 첫 항은 샘플이 많을수록 좋아지지만 두 번째 항은 N 이 커질수록 한없이 커진다. 이 두 번째 항이 바로 리워드 모델에 과하게 맞춰져 참된 리워드 기준 성능이 떨어지는 과최적화 오차다.

더 중요한 것은 하한도 함께 증명했다는 점이다. Best-of-N 을 쓰는 한 이 과최적화 항만큼의 손해는 항상 감수해야 한다는 뜻이므로, 샘플을 무작정 늘리면 성능이 떨어진다는 경험적 관찰이 이론으로 확정된 셈이다. 그 결과 Best-of-N 이 얻는 차수는 스카이라인보다 나쁘다. 상한(supremum)으로 정의되는 더 강한 유니폼 커버리지 조건을 가정하면 보장이 개선되지만, 보통 유니폼 커버리지가 일반 커버리지보다 훨씬 크기 때문에 여전히 스카이라인에는 닿지 못한다.

대안으로 제안된 알고리즘이 Inference-Time Pessimism 이다. 핵심은 참조 정책에서 너무 멀어지지 않도록 정규화한 목적함수의 최적 정책에서 표본을 뽑는 것인데, 이때 흔히 쓰는 KL 발산 대신 카이제곱 발산을 쓴다. 저자의 다른 연구에서 KL 정규화로는 스카이라인 차수에 도달할 수 없음이 밝혀졌기 때문이다. 실제로는 그 정책에서 직접 뽑을 수 없으므로, N 개의 표본으로 정규화 상수를 근사한 뒤 기각 표집(rejection sampling)으로 근사 표본을 얻고, N 번 안에 하나도 채택되지 않으면 Best-of-N 방식으로 되돌아간다. 정렬 과정 때문에 시간 복잡도가 N log N 으로 늘지만 성능 개선을 생각하면 감수할 만한 수준이라는 것이 발표자의 평가다.

주요 인사이트

• 이 논문의 가장 실용적인 메시지는 '샘플 수는 많을수록 좋다'는 통념이 Best-of-N 에서는 틀렸다는 것이다. 리워드 모델이 불완전한 이상, N 을 키우는 행위 자체가 그 오차를 파고드는 탐색이 된다.
• 제안 알고리즘의 후회 상한에는 N 항이 등장하지 않는다. 즉 일정 수준을 넘기면 N 이 얼마든 동일한 보장을 받으므로, 연산 예산을 늘렸다가 성능이 꺾일까 걱정할 필요가 없다.
• 정규화 방식의 선택이 이론적 최적성을 가른다는 점이 흥미롭다. RLHF 에서 관습적으로 쓰이던 KL 대신 카이제곱을 택한 것이 스카이라인 달성의 열쇠였다.
• 실무에서는 N 은 환경이 허락하는 만큼 크게 잡고 정규화 파라미터는 하이퍼파라미터 튜닝하듯 조정하라는 것이 발표에서 정리한 권고다.
• 일부 문제에서 Best-of-N 도 충분히 좋았다는 실험 결과는, 두 종류의 커버리지 값이 비슷한 경우라는 이론적 설명과 맞아떨어진다.

자주 묻는 질문

관련 AI 소식