큰 모델의 탐색 정책으로 작은 모델을 학습시켜 LLM 추론의 반복 샘플링을 넘는다
Beyond Repeated Sampling: Learning Search Policies for LLM Reasoning
무엇인가
이 논문은 LLM이 어려운 추론 문제를 풀 때 테스트 타임 컴퓨트를 늘리는 방식, 그중에서도 지배적인 전략인 순진한 반복 샘플링(naive repeated sampling)의 한계를 다룬다. 같은 모델에서 독립적인 답을 여러 개 뽑아 그중 하나가 맞기를 기대하는 방식인데, 이는 토큰 수준 디코딩 노이즈를 통해서만 탐색하기 때문에 서로 거의 중복되는 시도만 쏟아내고 진짜로 다른 아이디어를 만들어내지 못한다. 저자들은 탐색을 의미론적 수준에서 조종할 수 있는지, 즉 문제에 특화된 개념(concept), 힌트, 전략을 먼저 샘플링하고 그에 조건화해 답을 생성할 수 있는지 묻는다. 이는 어려운 문제를 풀 때 어떤 도구와 아이디어가 적용될지 먼저 고민하는 인간의 방식과 닮아 있고, RL 학습에서도 정답을 한 번도 발견하지 못하면 강화할 양성 신호 자체가 없다는 점에서 탐색 정책의 품질이 일급 관심사라고 주장한다.
어떻게 동작하나
논문의 출발점은 선행 개념 유도 샘플링(GuidedSampling, Handa et al. 2026)의 재현과 반박이다. 저자들은 MATH500에서 Qwen2.5-Instruct 1.5B, 3B, 7B, 14B, 32B를 개념 생성기(CG)와 답변 생성기(AG)로 각각 쓰는 25개 조합 전부를 평가하고, 문제당 총 100개 롤아웃을 배분해 pass@50을 측정했다. 원래 프로토콜은 반복 샘플링 베이스라인에 temperature 0.8, top-p 0.5라는 제한적 파라미터를 썼는데, 베이스라인과 개념 조건부 생성 모두를 탐색적 파라미터(temperature 1.0, top-p 0.95, top-k 비활성)로 바꾸면 개념 유도의 이득이 거의 모든 CG/AG 조합에서 사라지고 오히려 소폭 나빠진다. 개념 수 부족이 원인도 아니다. Llama-3.2-3B-Instruct는 문제당 3.18개의 개념을 생성해 Qwen보다 3배 많지만, 탐색적 샘플링에서는 86.7% 대 베이스라인 88.1%로 오히려 1.4포인트 뒤진다. 저자들은 또 두 가지 구조적 문제를 지적한다. 개념을 반복 루프에서 하나씩 생성하는 것은 자기회귀 모델에서 낭비이고, 문제당 평균 약 1개꼴로만 개념이 나와 다양성이라는 목적 자체가 무너진다는 것이다.
무엇과 다른가
이에 대한 1차 제안은 추론 시 절차의 개선이다. 개념을 하나씩 뽑는 대신, 개념 생성기에게 문제를 분석하고 유용한 개념을 단일 궤적(single trajectory) 안에서 한꺼번에 출력하도록 요청한다. 각 개념은 신호가 강하고 문제에 특화되어 있으며 중복되지 않아야 하며, 궤적에서 최대 10개까지 파싱해 답변 롤아웃을 조건화한다. 답변 생성은 temperature 1.0, top-p 0.95, top-k -1의 탐색적 설정을 유지해 비교를 보수적으로 만든다. 평가는 쉬운 문제를 배제하기 위해 답변 생성기 크기별로 반복 샘플링이 100개 롤아웃에서 0% 정확도를 얻은 문제만 남긴 난이도 부분집합에서 수행했고, 이 집합은 1.5B 모델의 493개에서 32B 모델의 178개로 줄어든다. 결과적으로 거의 모든 CG/AG 조합이 베이스라인을 앞서며 최대 +9.7 pass@50 포인트를 얻었고, 7B 생성기가 32B 답변 생성기를 +8.8포인트 끌어올렸다. 개념 수는 평균 4개에서 10개 사이로 늘었고 강한 생성기는 상한 10개를 채운다. Qwen 계열이 아닌 Llama-3.2-3B로도 424개 난이도 문제에서 25.2%에서 26.1%로(+0.9포인트) 소폭 개선되며 문제당 9.23개 개념을 생성했다.
어떻게 쓰나
2차 제안은 개념 생성기 자체를 학습 가능한 부품으로 만드는 것이다. 답변 생성기는 동결(frozen)한 채, 작은 개념 생성기 πθ가 더 큰 답변 생성기의 다운스트림 성공을 최대화하도록 강화학습으로 최적화한다. 문제마다 CG가 G=8개의 개념 궤적을 샘플링하고, 각 궤적에는 답변 생성기 호출 B=128회가 배정되며, 파싱된 개념 M개에 ⌊B/M⌋ 또는 ⌈B/M⌉개로 최대한 균등하게 분배한다(B=128, M=10이면 일부는 13회, 나머지는 12회). 각 롤아웃은 LLM 판정기로 정답 여부를 채점하고, 이를 궤적 단위 스칼라 보상으로 집계한다. 두 가지 집계 함수를 쓰는데, max-of-max는 궤적 내 어떤 개념에서 나온 어떤 롤아웃이든 맞으면 1을 주는 이진 보상이고, max-of-mean은 개념별 정확도의 최댓값으로 [0,1] 연속 보상이며 우연히 한 번 맞는 개념보다 안정적으로 유용한 개념을 보상한다. 보상은 개별 개념이 아니라 궤적 전체에 부여해 다양한 개념 집합을 탐색하도록 유도하고, 최적화는 OLMo3 변형의 GRPO 스타일 목적함수로 수행한다. 학습 데이터는 DeepMath-103k에서 답변 생성기 성공률 5% 미만 문제를 걸러 쓰고, 평가는 답변 생성기가 128개 롤아웃에서 0%를 기록한 1k 문제로 하며, 분포 밖 평가로 Omni-MATH 2를 같은 방식으로 필터링해 쓴다.
전제와 한계
주요 결과는 Qwen2.5-7B 개념 생성기와 동결된 Qwen2.5-32B 답변 생성기 조합에서 나온다. 가장 좋은 max-of-mean 변형은 pass@128 39.2%로 순진한 반복 샘플링의 19.0%를 약 두 배로 만들었고, pass@64에서도 29.6% 대 11.4%로 격차가 유지된다. 문제와 무관한 고정 힌트 10개를 쓰는 일반적 프롬프트 변형은 22.6%에 그쳤고, max-of-max는 35.3%였다. 학습하지 않은 개념 생성기와 비교하면 7B는 28.9%, 32B는 33.8%였는데, 학습된 7B가 자기보다 네 배 이상 큰 답변 생성기를 조종하면서 학습되지 않은 32B 생성기까지 앞선다. RL 학습은 학습되지 않은 7B 대비 +10.3포인트를 더했고, 분포 밖 Omni-MATH 2에서도 max-of-mean이 18.6% 대 순진한 샘플링 11.3%로 가장 좋았다(다만 max-of-max는 이 데이터셋에서 학습되지 않은 생성기와 동등한 수준이다).
분석 파트에서 저자들은 이득이 문제 특화 개념에서 온다는 것을 확인한다. 학습된 개념을 다른 문제에 무작위로 짝지어 주는 뒤섞기(derangement) 실험에서 불일치 개념은 23.9%로 일반 프롬프트 변형(22.6%) 수준에 그쳤고, 일치 개념 39.2%에 크게 못 미쳤다. 즉 순진한 샘플링 대비 약 20포인트 개선 중 약 15포인트가 문제 관련성에서 나온다. 정답 누출은 개념의 0.44%만이 정답을 포함하는 것으로 표시됐고 문제 기준 3%만 해당해 이득을 설명하지 못한다. 총 롤아웃 예산 B=128을 고정한 채 개념 수를 1개에서 10개로 늘리면 pass@128이 22.7%에서 39.8%로 오르지만 pass@1은 거의 변하지 않아, 고정 예산을 더 많은 개념에 분산하는 것이 해 공간 커버리지를 넓힌다는 것을 보여준다. 전이 실험에서는 Qwen2.5-32B에만 학습된 7B CG가 한 번도 본 적 없는 Llama-3.3-70B를 조종해 DeepMath에서 pass@128 34.3%를 기록, 순진한 샘플링 26.2%와 Llama가 스스로 만든 개념 28.9%를 모두 앞섰다. 같은 7B를 개념 생성기가 아니라 직접 답변하도록 RL 학습시키면 성능이 더 나빠, 이득은 작은 모델이 문제를 푸는 능력이 아니라 더 강한 모델을 조종하는 데서 온다. 계산 비용 측면에서 학습 스텝 약 413초 중 답변 롤아웃 300초와 판정 85초가 약 94%를 차지하고 개념 생성은 18초, 액터 업데이트는 10초에 불과하다. 추론 시에는 CG 호출 한 번이 답변 롤아웃 한 번의 약 0.24배(7.6B/32.5B × 1080/1060 토큰)로, B=128 배정의 0.3% 미만이다.
개발자 관점에서 이 논문이 주는 실무적 패턴은 명확하다. 큰 모델을 직접 파인튜닝하기 어렵거나 가중치가 공개되지 않은 API 모델만 쓸 때, 작은 모델을 그 큰 모델의 검색 정책으로 한 번 학습시켜 두면 추론 시 거의 추가 비용 없이 pass@k를 끌어올릴 수 있다. 다만 도입 전에 확인할 것이 있다. 첫째, 이득은 반복 샘플링 베이스라인을 탐색적 디코딩 파라미터로 공정하게 돌렸을 때만 의미가 있으므로, 자체 평가에서도 temperature와 top-p를 충분히 탐색적으로 설정해야 한다. 둘째, 이득은 모델이 100~128개 롤아웃으로도 전혀 못 푸는 난이도 구간에서 측정된 것이므로 쉬운 문제가 많은 워크로드에서는 효과가 작을 수 있다. 셋째, 개념을 문제별로 최대 10개 파싱해 롤아웃을 균등 분배하는 구조라 파싱 실패나 개념 수 변동에 대한 견고성을 자체 점검해야 한다. 넷째, 보상 추정에 답변 생성기 롤아웃과 LLM 판정이 대량 필요해 학습 자체는 계산 집약적이며, 학습은 한 번만 하면 된다는 전제가 있다.
저자들이 밝힌 한계도 분명하다. 학습이 계산 집약적인데 그 비용은 거의 전부 보상 추정, 즉 학습하지 않는 큰 답변 생성기를 반복 실행하는 데서 발생한다. 보상이 개별 개념이 아니라 궤적 단위로 부여되어 약한 개념이 강한 개념과 신용을 공유할 수 있고 위치 효과 가능성도 남는다. 평가는 수학적 추론에 집중되어 있어 다른 도메인으로의 일반화는 검증되지 않았다. 향후 과제로 더 낮은 계산의 보상 추정, 더 세밀한 신용 할당, 판정기에 견고한 평가, 개념 간 롤아웃 배분을 위한 밴딧 기반 접근, 그리고 더 넓은 도메인으로의 확장을 제시한다.