과정 보상 모델(PRM) 보정과 문제 난이도별 추론 계산량 조절, 논문 리뷰로 본 원리

핵심 메시지

• 요즘 상용 AI 서비스는 질문 난이도를 스스로 평가해 얼마나 오래 생각할지 정하지만, 이를 정면으로 다룬 학술 연구는 아직 많지 않다.
• 베스트 오브 N이나 빔 서치 같은 기존 추론 단계 스케일링 기법은 문제 난이도와 무관하게 계산량이 고정돼 있다.
• 과정 보상 모델(PRM)의 점수를 0과 1 사이로 정규화해 '지금까지의 생성으로부터 끝까지 갔을 때 정답일 확률'로 해석하면 필요한 샘플 수를 수식으로 계산할 수 있다.
• 그런데 기존 PRM은 실제 성공 확률보다 점수를 후하게 주는 경향이 있어, 분위수 회귀로 보정해야 이 계산이 성립한다.
• 보정한 PRM에 문제별 적응형 스케일링을 붙이면 기존 방식과 비슷하거나 더 나은 정확도를 훨씬 적은 계산 예산으로 얻는다.

쉽게 이해하기

같은 챗봇에 질문을 던져도 어떤 질문에는 즉답이 오고 어떤 질문에는 한참 뜸을 들인다. 최근 모델들은 질문이 들어오면 난이도를 내부적으로 평가해 어떤 모델로 답할지, 또는 같은 모델이라도 얼마나 오래 생각한 뒤 답할지를 스스로 정한다. 발표자는 이런 기능이 서비스에는 이미 들어와 있지만, 추론 단계의 계산량을 문제에 맞춰 적응적으로 조절하는 방법을 정면으로 다룬 논문은 의외로 드물다고 짚으며 이 논문을 골랐다.

출발점은 추론 단계 스케일링이다. 모델을 추가로 학습시키지 않고, 이미 학습된 모델에 추론 시점의 계산량만 더 써서 성능을 끌어올리는 방식이다. 대표적인 것이 답변 후보를 N개 만들어 점수가 가장 높은 것을 고르는 베스트 오브 N, 그리고 중간중간 점수를 매겨 상위 후보만 남기며 이어 쓰는 빔 서치다. 문제는 두 방법 모두 N이나 후보 개수가 고정돼 있다는 점이다. 쉬운 문제에 과하게 계산을 쓰면 전력과 비용이 낭비되고, 어려운 문제를 성급히 끝내면 틀릴 확률이 커진다.

논문의 첫 번째 아이디어는 과정 보상 모델(PRM)의 점수를 확률로 다시 읽는 것이다. 질문과 지금까지 생성된 부분 답변을 합친 상태에서 모델이 끝까지 생성했을 때 정답을 맞힐 확률을 안다면, 목표 성공률을 넘기기 위해 몇 번을 뽑아야 하는지는 간단한 확률 계산으로 나온다. 주사위를 한 번 던져 1이 나올 확률은 낮지만 여러 번 던지면 한 번이라도 1이 나올 확률이 올라가는 것과 같은 원리다. 이렇게 구한 최소 샘플 수를 문제마다 다르게 적용하되, 아무리 어려워도 무한정 붙잡고 있지 않도록 계산 예산의 상한을 함께 둔다.

문제는 기존 PRM 점수를 그냥 0과 1 사이로 정규화해도 실제 성공 확률과 잘 맞지 않는다는 점이다. 저자들의 측정에서 기존 PRM은 실제보다 점수를 후하게 주는 쪽으로 치우쳐 있었다. 그래서 두 번째 아이디어로 분위수 회귀를 이용한 보정이 등장한다. 오차가 좌우 대칭이라는 가정이 깨진 상황에서는 평균을 맞추는 일반적인 회귀 대신 특정 분위수를 맞추는 손실 함수를 쓰는 편이 낫고, PRM의 출력 차원을 여러 개로 늘려 여러 분위수 값을 한꺼번에 내놓게 학습시킨다. 학습 데이터는 같은 지점에서 여러 번 끝까지 생성해 정답 비율을 세는 방식으로 만든다.

실험에서 보정한 PRM은 여러 지표에서 실제 성공 확률에 더 가깝게 맞았다. 보정 없이 적응형 스케일링을 하면 사실상 추론 스케일링을 안 한 것과 다름없는 결과가 나온 반면, 보정 뒤에는 기존 방식과 비슷하거나 더 나은 정확도를 절반 이하의 예산으로 냈다. 발표자는 아이디어가 이론적으로도 깔끔하다고 평가하면서 한계도 짚는다. 추론 단계 스케일링의 매력은 추가 학습 없이 성능을 올리는 데 있는데 이 방법은 PRM을 따로 학습시켜야 한다. 보정에 쓰는 성공 확률을 실제 사용하는 언어 모델로 추정하기 때문에, 쓰는 모델이 바뀌면 PRM도 다시 학습시켜야 한다.

주요 인사이트

• '얼마나 생각할지'를 정하는 문제는 결국 확률 추정 문제다. 현재 상태에서의 성공 확률을 알면 필요한 시도 횟수는 계산으로 나온다.
• 보상 모델의 점수는 그 자체로 확률이 아니다. 점수를 확률처럼 쓰려면 실제 성공률과 맞추는 보정 과정이 반드시 필요하다.
• 보정을 건너뛴 적응형 스케일링은 계산만 아끼고 성능은 얻지 못한다. 절약처럼 보이는 결과가 사실은 아무것도 하지 않은 것과 같을 수 있다.
• 분위수를 여러 개 학습시키면 정확도와 비용 사이의 트레이드오프를 운영 단계에서 조절할 수 있는 손잡이가 된다.
• 사용하는 언어 모델이 바뀔 때마다 보상 모델을 다시 맞춰야 한다는 점은 이 방식의 실용성을 가늠하는 현실적인 기준이다.

자주 묻는 질문

관련 AI 소식