💬 NLP

Reuse your FLOPs: Scaling RL on Hard Problems by Conditioning on Very Off-Policy Prefixes

PrefixRL은 성공적인 궤적의 오프-폴리시 접두사(off-policy prefixes)에 조건화함으로써 학습을 안정화하고 샘플 효율성을 높여 어려운 추론 문제에서의 강화 학습의 비효율성을 해결하며, 표준 베이스라인 대비 더 빠른 수렴과 우수한 성능을 달성하는 동시에 거절 샘플링(rejection sampling)을 통한 자기 개선 루프를 가능하게 합니다.

원저자: Amrith Setlur, Zijian Wang, Andrew Cohen, Paria Rashidinejad, Sang Michael Xie

게시일 2026-02-04

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Amrith Setlur, Zijian Wang, Andrew Cohen, Paria Rashidinejad, Sang Michael Xie

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 학생(AI 모델)에게 매우 어려운 수학 문제를 푸는 법을 가르치고 있다고 상상해 보세요. 그 학생은 똑똑하지만, 정말 어려운 문제에 직면하면 그냥 추측해서 풀다가 틀리곤 합니다. 사실, 너무 자주 틀리기 때문에 학습 자체가 멈춰버립니다. 왜냐하면 공부할 만한 '정답'을 한 번도 마주하지 못하기 때문입니다. 이것이 바로 이 논문이 다루는 문제입니다: 강화 학습(RL)은 AI가 승리하는 움직임을 전혀 보지 못할 때 매우 어려운 문제에서 정체되는 경향이 있습니다.

저자들은 PrefixRL이라는 영리하고 새로운 방법을 제안합니다. 이 방법의 작동 원리를 쉬운 비유를 통해 설명하겠습니다.

1. 문제점: "어둠 속에 갇힌" 시나리오

학생이 출구를 찾기 위해 어두운 미로 속에 있다고 상상해 보세요. 발걸음을 옮길 때마다 벽에 부딪히고 "0점"이라는 신호를 받습니다. 학생은 계속 제자리를 맴돌며 에너지(컴퓨터 자원)를 낭비하지만, 출구로 가는 길은 결코 찾지 못합니다. 출구를 발견하지 못하기 때문에, 어느 방향이 맞는지 알 수 없으며 학습은 중단됩니다.

AI 용어로 말하자면, 모델이 어려운 문제를 해결하려고 노력하지만 스스로 정답을 생성하는 경우가 거의 없을 때 이런 현상이 발생합니다. '보상'(당신이 해냈다는 신호!)이 너무 드물어서 AI가 길을 잃게 되는 것입니다.

2. 기존 방식: "그냥 답을 암기하기"

이전에는 연구자들이 이전 시도에서 얻은 몇 가지 정답(off-policy 데이터)이 있다면, RL을 다시 시도하기 전에 학생에게 먼저 그 답들을 암기하도록 강요하곤 했습니다(지도 미세 조정, Supervised Fine-Tuning).

결함: 이것은 학생에게 미로를 통과하는 특정 경로를 암기하도록 강요하는 것과 같습니다. 일단 암기하고 나면, 학생은 탐색을 멈춥니다. 미로가 약간 변하거나 더 나은 경로를 찾아야 할 때, 학생은 호기심과 창의성을 잃었기 때문에 갇혀버리게 됩니다. 즉, 너무 경직되게 됩니다.

3. 새로운 방식: PrefixRL (The "Head Start" Strategy)

PrefixRL은 게임의 판도를 바꿉니다. 연구자들은 학생에게 전체 답을 암기하게 하는 대신, **'헤드 스타트(선행 출발)'**를 제공합니다.

비유: 학생이 다시 어두운 미로에 갇혔다고 상상해 보세요. 이번에는 친구(미로를 이미 풀어본 사람)가 다음과 같이 적힌 종이를 건넵니다: "너는 지금 빨간 문이 있는 모퉁이에 있어. 여기서부터는 왼쪽으로 돌아."
학생은 빨간 문까지 가는 법을 알아낼 필요가 없습니다. 그저 거기서부터 시작하면 됩니다.
핵심: 학생은 여전히 나머지 미로를 스스로 풀어내야 할 책임이 있습니다. 단순히 전체 경로를 베끼는 것이 아니라, 그 '헤드 스타트'를 이용해 여정의 '나머지 부분'을 연습하는 것입니다.

기술적으로 말하면, AI는 과거에 발견된 정답을 가져와서 그 앞부분(prefix)을 잘라낸 뒤, 이를 문제에 붙입니다. 그런 다음 AI는 남은 부분을 완성하려고 시도합니다. '좋은 지점'에서 시작하기 때문에, AI는 보상을 받을 확률이 훨씬 높아지며 이를 통해 학습할 수 있습니다.

4. 마법 같은 기술: "역방향 일반화 (Back-Generalization)"

이 논문에서 가장 놀라운 발견은 저자들이 **'역방향 일반화(Back-Generalization)'**라고 부르는 것입니다.

비유: 당신이 특정하고 쉬운 고속도로 구간에서만 운전 연습을 한다고 상상해 보세요. 당신은 그 구간에서 합류하고, 가속하고, 조향하는 법을 완벽하게 배웁니다.
놀라운 점: 비록 당신이 험난하고 구불구불한 산길(원래의, prefix가 없는 문제)에서는 한 번도 연습해 본 적이 없지만, 고속도로에서의 운전 실력이 어떻게 당신을 산길에서도 더 나은 운전자로 만들어 주었을까요?
이유: 논문은 AI가 '헤드 스타트' 문제가 가진 논리와 전략을 연습함으로써, 문제의 근저에 깔린 논리와 전략을 배우게 된다고 설명합니다. 즉, 무엇을 생각할지가 아니라 어떻게 생각할지를 배우는 것입니다. 따라서 헤드 스타트 없이 원래의 어려운 문제로 돌아갔을 때, AI는 그 새로운 전략들을 적용하여 이전보다 더 잘 해결할 수 있게 됩니다.

5. 왜 더 나은가 (결과)

논문은 이 방법을 매우 어려운 수학 및 코딩 문제에 테스트했습니다.

속도: PrefixRL은 기존의 가장 좋은 방법들보다 두 배 더 빠르게 학습했습니다. 계속 어둠 속에서 헤매지 않아도 되기 때문에 컴퓨터 자원을 덜 낭비했습니다.
품질: 최종적인 AI는 기존 방식보다 어려운 문제를 해결하는 능력이 3배 더 뛰어났습니다.
유연성: 심지어 '헤드 스타트'가 완전히 다른 유형의 AI(예: Qwen 모델의 힌트를 사용하여 Llama 모델을 훈련하는 경우)로부터 왔을 때도 효과적이었습니다. 이는 똑같은 AI가 힌트를 생성할 필요는 없으며, 어떤 똑똑한 AI라도 '헤드 스타트'를 제공할 수 있음을 의미합니다.

요약

PrefixRL은 어려움에 부딪힌 학생에게 질문의 가장 어려운 부분을 통과할 수 있는 힌트를 주어, 나머지 부분을 풀며 연습할 수 있게 하는 것과 같습니다. 놀랍게도, 힌트를 가지고 연습함으로써 학생은 문제의 '논리' 자체를 아주 잘 익히게 되어, 힌트가 없는 원래의 버전조차 이전보다 더 잘 풀 수 있게 됩니다. 이는 AI에게 단순히 답을 암기하도록 강요하지 않으면서도, 기존의 컴퓨터 작업량을 재사용하여 학습을 더 빠르고 똑똑하게 만드는 방법입니다.

기술 요약: 당신의 FLOPs를 재사용하라: 매우 오프-폴리시(Very Off-Policy) 접두사 조건화를 통한 어려운 문제에 대한 RL 스케일링

1. 문제 정의

강화 학습(RL)은 대규모 언어 모델(LLM)의 추론 능력을 향상시키는 데 있어, 특히 수학 및 코딩 분야에서 사실상의 표준 방법론이 되었습니다. 그러나 표준 온-폴리시(On-policy) RL 방식은 기본 모델의 통과율(pass rate)이 거의 0에 가까운 어려운 문제(예: pass@2k ≈ 0)에 적용될 때 심각한 병목 현상에 직면합니다.

이러한 체제에서 RL 정책은 올바른 추론 궤적(reasoning trace)을 샘플링하는 경우가 매우 드뭅니다. 결과적으로 다음과 같은 문제가 발생합니다:

학습 신호 소멸: 샘플링된 모든 궤적이 보상(reward)을 0으로 받기 때문에 정책 그래디언트 추정치가 0 또는 0에 가깝게 됩니다.
컴퓨팅 낭비: 시스템은 학습 신호를 전혀 받지 못한 채 샘ming에 방대한 양의 FLOPs를 소비하며, 이로 인해 학습이 정체되고 보상이 정체(plateau)됩니다.
기존 솔루션의 비효 효율성:
- 오프-폴리시 궤적에 대한 지도 미세 조정(SFT): 이는 모델을 부트스트래핑할 수 있지만, 종종 엔트로피 붕괴(entropy collapse)와 암기(memorization)를 초래하여 후속 탐색을 저해합니다.
- 직접적인 오프-폴리시 RL: 중요도 가중치(importance weighting)를 통해 오프-폴리시 궤적을 직접 사용하는 것은 높은 분산의 그래디언트 추정과 불안정성을 야기합니다.

핵적인 과제는 (드물게 나타나는 정답 궤적을 포함할 수 있는) 이전의 추론 또는 RL 실행으로부터 얻은 "낭비된" 컴퓨팅 자원을, 직접적인 오프-폴리시 감독과 관련된 불안정성을 도입하지 않으면서 어떻게 온라인 RL 정책을 가이드하는 데 재사용할 것인가 하는 점입니다.

2. 방법론: PrefixRL

저자들은 PrefixRL을 제안합니다. 이는 오프-폴리시 궤적을 직접적인 감독 대상으로 사용하는 대신, 성공적인 오프-폴리시 궤적의 접두사(prefix)에 조건을 거는(conditioning) 온-폴리시 RL 프로세스입니다.

핵심 메커니즘

데이터 수집: 각 문제에 대해 정답 궤적이 발견될 때까지 기본 모델에 대한 대규모 거절 샘플링(rejection sampling)을 수행하여 생성된 오프-폴리시 궤적 데이터셋( $D_{off}$ )을 수집합니다.
문제 증강: 각 원래의 문제 $x$ 에 대해, 시스템은 올바른 오프-폴리시 궤적 $y_x$ 로부터 접두사(첫 $h$ 개의 토큰)를 추출합니다. 이 접두사는 원래의 문제 프롬프트에 결합되어 **접두사가 붙은 문제(prefixed problem)**를 생성합니다: $\text{concat}(x, y_{x, 1:h})$ .
학습 목적 함수: 모델은 다음의 혼합 데이터로 온-폴리시 RL 학습을 진행합니다:
- 원래의 (접두사가 없는) 문제들.
- 접두사가 붙은 문제들.
- 결정적으로: 그래디언트는 오프-폴리시 접두사 토큰에 대해 마스킹됩니다. 모델은 오직 접두사 이후에 생성하는 토큰에 대해서만 업데이트됩니다.

이론적 보장

본 논문은 PrefixRL 목적 함수에 대한 이론적 증명을 제공합니다:

일관성(Consistency): 만약 오프-폴리시 접두사가 올바르고 모델 클래스 내의 정책에 의해 실현 가능하다면, PrefixRL 목적 함수의 극대화(maximizer)는 표준 RL 목적 함수의 극대화이기도 합니다. 즉, 이 방법은 최적해에 편향을 주지 않습니다.
샘플 효율성(Sample Efficiency): PrefixRL은 표준 RL에 비해 하위 최적성 격차(suboptimality gap)를 줄여줍니다. 정책을 더 높은 보상을 받는 상태(정답 솔루션이 이미 부분적으로 형성된 상태)에 배치함으로써 탐색 부담을 줄입니다. 이론적 경계값은 필요한 온-폴리시 샘플 수가 컨텍스트 길이에 따라 유리하게 스케일링되며 감소함을 시사합니다.

주요 현상: 역방향 일반화(Back-Generalization)

경험적으로 저자들은 역방향 일반화를 발견했습니다. 즉, 접두사가 붙은 문제로만 학습하더라도, 모델이 해당 단계에서 접두사가 없는 버전에 대해 명시적으로 학습하지 않았음에도 불구하고 원래의 (접두사가 없는) 문제에 대한 성능이 크게 향과됩니다.

메커니즘: 모델은 단순히 접두사를 모방하는 것이 아닙니다. 대신, 접두사로부터 전략을 학습합니다.
전략 발견: 모델은 접두사에 존재하는 하위 최적의 전략을 학습하지 않을 수 있으며(unlearn), 새로운 더 나은 전략을 발견할 수 있습니다. 모델은 보이지 않는 상태(unseen states)에 대한 다음 토큰 분포를 변경함으로써, 접두사가 제공하는 고보상 상태로부터 역추적(backtracking)하여 더 나은 경로를 찾아냅니다.
관련성: 역방향 일반화는 접두사(또는 인-컨텍스트 문제)가 타겟 문제와 의미론적으로 관련이 있을 때 가장 효과적입니다.

3. 주요 기여

PrefixRL 프레임워크: 올바른 접두사에 조건을 걸어 오프-폴리시 FLOPs를 재활용하는 새로운 방법으로, 직접적인 오프-폴리시 RL이나 SFT로 인한 엔트로피 붕괴를 피합니다.
이론적 분석: PrefixRL이 표준 RL 목적 함수와 일관되며, 탐색 부담을 줄임으로써 증명 가능한 샘플 효율성 이점을 제공한다는 것을 입증했습니다.
역방향 일반화의 발견: 접두사가 붙은 문제에 대한 학습이 원래의 문제에 대한 성능을 높이는 현상을 확인하였으며, 이는 단순한 일반화를 넘어 관련 작업 전반에 걸쳐 우수한 성능을 보이는 강력한 메커니즘임을 보여줍니다.
자기 개선 루프(Self-Improvement Loop): 오프-폴리시 궤적이 외부 교사 모델 없이 기본 모델 자체의 거절 샘플링을 통해 생성되는 실용적인 파이프라인을 시연하여, 지속적인 자기 개선 사이클을 구축했습니다.

4. 실험 결과

저자들은 Llama-3.1-8B 및 Qwen-3-4B 모델을 사용하여 어려운 수학 추론 벤치마크(DAPO, OMNI-MATH, AIME '25, HMMT '25)에서 PrefixRL을 평가했습니다.

컴퓨팅 효율성: PrefixRL은 초기 거절 샘플링의 컴퓨팅 비용을 고려하더라도, 가장 강력한 베이스라인(오프-폴리시 데이터에 대한 SFT 후 RL 수행)보다 2배 더 빠르게 동일한 훈련 보상을 달단합니다.
정확도 향상:
- 어려운 훈련 문제에 대해, PrefixRL은 베이스라인 대비 최종 정확도를 >45%(상대적으로 3배 이상) 향상시킵니다.
- AIME '25와 같은 홀드아웃(held-out) 벤치마크에서, PrefixRL은 컴퓨팅 매칭 비교 시 가장 강력한 베이스라인보다 pass@1을 12% 개선합니다.
출처에 대한 강건성: 이 방법은 오프-폴리시 접두사가 다른 모델 패밀리에서 유래한 경우에도(예: Qwen 접두사를 사용하여 Llama 정책을 훈련) 효과적으로 작동하여, 실제 환경에서의 유연성을 입증합니다.
훈련 역학:
- 안정성: SFT+RL과 달리, PrefixRL은 토큰 수준의 엔트로피를 보존하여 탐색 붕괴를 방지합니다.
- 그래디언트 품질: PrefixRL은 표준 오프-폴리시 RL(높은 분산의 중요도 가중치로 고통받는)에 비해 더 높은 그래디언트 노름(norm)과 낮은 분산을 보입니다.
- 효율성: 모델은 평균적으로 더 짧은 응답을 생성하며, 이는 모델이 더 결정적인 단계에 더 빨리 도달하고 "생산적이지 않은 방황"을 줄임을 나타냅니다.

5. 의의 및 주장

본 논문은 PrefixRL이 어려운 문제에 대한 RL 스케일링 접근 방식을 근본적으로 변화시킨다고 주장합니다. 오프-폴리시 데이터를 모방의 대상(이는 불안정성이나 붕괴를 초래함)으로 보거나 노이즈로 버리는 대신, PrefixRL은 이를 탐색을 가이드하는 조건화 메커니즘으로 취급합니다.

이 연구의 의의는 다음과 같습니다:

FLOPs 재사용: 이전에 낭비되었던 컴퓨팅 자원(거절 샘플링을 통해 정답 궤적을 얻어낸 실패한 시도들)을 학습을 위한 직접적인 신호로 전환합니다.
"어려운 문제" 체제 해결: 표준 RL이 정체되는 영역, 즉 온-폴리시 롤아웃이 거의 성공하지 못하는 문제에서도 RL 훈련을 가능하게 합니다.
역방향 일반화: 접두사에 대한 조건화가 접두사가 없는 작업에 대해 우수한 성능을 이끌어낼 수 있다는 발견은, 단순한 일반화를 넘어 과거의 컴퓨팅 자원을 활용하는 강력한 메커니즘을 제시하며, 향후 RL 알고리즘이 역사적 데이터를 통합하는 방식을 재편할 수 있음을 시사합니다.

저자들은 역방향 일반화를 활용하는 것이 이전의 컴퓨팅 자원을 부트스트래핑하고 추론 LLM의 지속적인 자기 개선을 가능하게 하는 핵심 요소라고 결론짓습니다.