WARP: 정책을 weight space에서 세 번 병합한다

WARP: On the Benefits of Weight Averaged Rewarded Policies (Ramé et al., Google DeepMind, arXiv 2024)

Introduction

#19 PPO에서 봤듯, RLHF의 policy 학습은 reward 최적화와 KL 정규화 사이의 줄다리기다. reward model이 매긴 점수를 최대화하는 방향으로 정책을 밀어붙이되, 그 정책이 SFT 초기화로부터 너무 멀어지지 않도록 KL divergence 페널티로 붙잡아 둔다. 이 KL 항이 없으면 정책은 #10 reward model overoptimization에서 다룬 것처럼 proxy reward만 파고들다 실제 응답 품질을 잃어버린다.

문제는 이 KL 정규화가 양날의 검이라는 점이다. WARP 논문은 이 트레이드오프를 다음과 같이 정리한다: KL 정규화는 “forces the policy to remain close to its supervised fine-tuned initialization, though it hinders the reward optimization” — 정책을 SFT 초기화 근처에 묶어 두어 망각과 reward hacking을 막지만, 동시에 reward 최적화 자체를 방해한다는 것이다. KL 계수 \(\beta\)를 키우면 정책이 안전해지는 대신 reward가 덜 오르고, \(\beta\)를 줄이면 reward는 오르지만 망각과 hacking 위험이 커진다. 이 트레이드오프는 스칼라 하나로는 근본적으로 풀리지 않는다.

WARP(Weight Averaged Rewarded Policies)는 이 줄다리기를 KL 계수 튜닝이 아니라 weight space에서의 병합으로 우회한다. 세 단계로 구성된다.

  1. EMA 앵커: 고정된 SFT 초기화 대신 정책 자신의 exponential moving average를 KL 정규화의 동적 앵커로 쓴다.
  2. SLERP 병합: 독립적으로 학습한 \(M\)개 정책의 task vector를 구면 선형 보간으로 합친다.
  3. LITI: 병합된 모델을 다시 초기화 쪽으로 선형 보간해 사전학습 특성을 되찾는다.

이 세 단계를 반복하면서 KL-reward 파레토 프런티어 자체를 점진적으로 밀어올린다. 흥미로운 대목은 이 논문이 #13 WARM에 대한 응답으로 쓰였다는 점이다. WARM이 여러 reward model을 평균해 reward 설계 쪽에 weight averaging을 적용했다면, WARP는 같은 아이디어를 정책 학습 쪽에 적용한다.

Background

KL 정규화의 원형

#19 PPO의 RLHF objective는 다음과 같은 형태다.

\[\max_{\theta} \; \mathbb{E}_{x, y \sim \pi_\theta} \left[ r(x, y) \right] - \beta \, \mathrm{KL}\left[\pi_\theta(\cdot \mid x) \,\|\, \pi_{init}(\cdot \mid x)\right]\]
  • \(r(x, y)\): reward model이 매긴 점수
  • \(\pi_{init}\): SFT로 얻은 고정된 초기 정책
  • \(\beta\): KL 페널티의 세기

여기서 앵커 \(\pi_{init}\)는 학습 내내 고정이다. WARP가 건드리는 지점이 바로 이 앵커다.

Overoptimization과 파레토 프런티어

#10 reward model overoptimization에서 확인한 것처럼, KL divergence를 가로축에 proxy reward를 세로축에 두면 초반에는 KL이 늘어날수록 reward도 따라 오르지만, 어느 지점을 넘어서면 proxy reward는 계속 오르는데 gold reward(진짜 사람이 평가하는 품질)는 꺾인다. 즉 “KL 예산 안에서 얼마나 높은 진짜 reward를 얻을 수 있는가”를 나타내는 파레토 프런티어가 존재하고, 이 프런티어 자체를 밀어올리는 것이 개선의 방향이다. WARP의 주장은 weight space 병합이 정확히 이 프런티어를 이동시킨다는 것이다.

Model soups와 WARM

여러 모델의 가중치를 단순 평균하는 model soups는 fine-tuning 커뮤니티에서 이미 널리 쓰인다. #13 WARM은 이 아이디어를 reward model 여러 개를 평균하는 데 적용해, 개별 RM보다 reward hacking에 강건한 단일 RM을 만들었다. WARP는 같은 weight averaging 아이디어를 쓰지만 두 가지가 다르다. 첫째, 평균 대상이 RM이 아니라 정책이다. 둘째, 표준 선형 평균(LERP) 대신 task vector의 norm을 보존하는 SLERP를 쓰고, 이를 RL 정렬 맥락에서 반복 적용한다.

구분 WARM (#13) WARP
병합 대상 여러 reward model 가중치 여러 정책(policy) 가중치
파이프라인 위치 reward 설계 단계 정책 학습 단계
병합 방식 linear averaging SLERP + LITI, 반복 적용
목적 reward hacking에 강건한 RM KL-reward 파레토 프런티어 이동

Method

Stage 1 — EMA 앵커

고정된 \(\pi_{init}\) 대신, 정책 자신의 exponential moving average를 KL 정규화 앵커로 쓴다.

\[\theta_{ema} \leftarrow (1-\mu)\,\theta_{ema} + \mu\,\theta_{policy}\]
  • \(\theta_{policy}\): 매 스텝 업데이트되는 정책 가중치
  • \(\theta_{ema}\): KL 앵커로 쓰이는 이동평균, \(\mu = 0.01\)로 아주 천천히 갱신된다

이 앵커는 mean teacher 방식의 증류로도 볼 수 있다. 효과는 “automatic annealing of the KL regularization” — 학습 초반에는 \(\theta_{ema}\)가 초기화에 가까워 KL 페널티가 강하게 걸리지만, 학습이 진행되며 \(\theta_{ema}\)가 \(\theta_{policy}\)를 서서히 뒤쫓아가면서 두 가중치 사이의 거리가 좁아지고 실질적인 제약이 자연히 완화된다. \(\beta\) 하나를 고정해 놓고 튜닝하는 대신, 학습 진행 상황에 맞춰 제약의 세기가 스스로 조절되는 셈이다.

실제로 \(\beta = 0.1\)에서 EMA 앵커는 여러 KL 수준에 걸쳐 고정 앵커를 쓰는 변형(\(\beta = 0.01\), \(\beta = 0.001\))보다 높은 reward를 얻었다. 고정 앵커는 KL 예산을 늘리려면 \(\beta\)를 계속 낮춰야 하지만, 이는 학습 전 구간에 걸쳐 균일하게 제약을 풀어버려 초반의 안전장치까지 함께 잃는다. 동적 앵커는 초반엔 강하게 묶고 후반에만 풀어주는 암묵적 annealing을 만들어 이 문제를 피한다.

Stage 2 — SLERP 병합

\(I\)번째 iteration에서 동일한 초기화 \(\theta_{init}\)로부터 \(M\)개 정책을 독립적으로 학습한다. 각 정책의 task vector는 \(\delta^m = \theta^m - \theta_{init}\)이다. 이 task vector들을 구면 선형 보간(spherical linear interpolation)으로 합친다. \(M=2\)일 때는 다음과 같다.

\[\mathrm{slerp}(\theta_{init}, \theta^1, \theta^2, \lambda) = \theta_{init} + \frac{\sin((1-\lambda)\Omega)}{\sin \Omega}\delta^1 + \frac{\sin(\lambda\Omega)}{\sin \Omega}\delta^2\]
  • \(\Omega\): 두 task vector \(\delta^1\), \(\delta^2\) 사이의 각도
  • \(\lambda\): 보간 비율, 두 정책을 대칭적으로 섞을 때는 \(\lambda = 0.5\)

이 연산은 레이어 단위로 적용한다. 표준 선형 평균(LERP)은 \(\delta^1\)과 \(\delta^2\)가 서로 다른 방향을 가리킬수록 평균 벡터의 norm이 줄어드는 경향이 있는데(서로 다른 방향의 벡터가 상쇄되기 때문), SLERP는 구면을 따라 보간하기 때문에 이 norm 손실을 줄인다. 그 결과 SLERP로 합친 모델은 개별 정책보다 높은 reward를 얻는 대신 KL이 약간 늘어난다 — norm이 보존된다는 것은 곧 초기화로부터 더 멀리 이동했다는 뜻이기도 하다.

Stage 3 — LITI

SLERP로 합친 모델 \(\theta_{slerp}\)는 reward는 높지만 그만큼 초기화에서 멀어져 사전학습 때 가지고 있던 특성(문법, 안전성, 다양성)을 일부 잃는다. LITI(Linear Interpolation Toward Initialization)는 이 모델을 다시 초기화 쪽으로 선형 보간한다.

\[\theta^\eta \leftarrow (1-\eta)\,\theta_{init} + \eta\,\theta_{slerp}\]

\(\eta \in [0,1]\)을 0에서 1까지 조절하면서 매번 평가하면, \(\eta = 0\)(순수 초기화)부터 \(\eta = 1\)(순수 SLERP 병합 모델)까지 KL과 reward가 함께 움직이는 궤적이 드러난다. 이 궤적이 개별 정책 하나로는 얻을 수 없는 개선된 파레토 프런티어를 이룬다. 논문은 보통 \(\eta = 0.3\) 부근을 골라 다음 iteration의 초기화로 삼는다.

토이 예제: 2차원 weight space에서 손으로 따라가기

세 단계가 실제로 reward와 KL을 어떻게 움직이는지, 아주 작은 2차원 예시로 직접 계산해본다. (아래 수치는 논문 수치가 아니라 이해를 돕기 위해 구성한 예시다.)

초기화를 원점에 둔다: \(\theta_{init} = (0, 0)\). 두 정책을 독립적으로 학습시켰더니 서로 직교하는 방향으로 움직였다고 하자.

\[\delta^1 = (4, 0), \qquad \delta^2 = (0, 3)\]

두 벡터의 내적이 0이므로 \(\Omega = 90°\)다. \(\lambda = 0.5\)로 SLERP를 계산하면

\[\frac{\sin(45°)}{\sin(90°)} \approx 0.707\]

이므로

\[\theta_{slerp} = 0.707 \cdot (4,0) + 0.707 \cdot (0,3) \approx (2.83,\, 2.12)\]

이 벡터의 norm은 \(\sqrt{2.83^2 + 2.12^2} \approx 3.54\)다. 반면 같은 두 벡터를 그냥 평균(LERP)하면

\[\theta_{lerp} = 0.5\cdot(4,0) + 0.5\cdot(0,3) = (2,\, 1.5), \qquad \|\theta_{lerp}\| = 2.5\]

개별 task vector의 norm은 각각 4와 3이고 그 평균은 3.5다. SLERP 결과의 norm(3.54)은 이 평균에 가깝게 보존되는 반면, LERP 결과의 norm(2.5)은 두 벡터가 서로 다른 방향을 가리키기 때문에 상쇄되어 줄어든다. norm이 크다는 것은 초기화로부터 더 멀리, 즉 reward가 높아지는 방향으로 더 많이 이동했다는 뜻이므로 SLERP가 LERP보다 높은 reward를 보존하는 이유가 여기서 드러난다.

이제 LITI로 \(\theta_{slerp}\)를 초기화 쪽으로 당겨본다. \(\eta\)를 바꿔가며 결과 norm을 보면 다음과 같다.

\(\eta\) \(\theta^\eta\) norm
0.0 \((0,\, 0)\) 0
0.3 \((0.85,\, 0.64)\) 1.06
0.6 \((1.70,\, 1.27)\) 2.12
1.0 \((2.83,\, 2.12)\) 3.54

\(\eta\)가 커질수록 초기화로부터 멀어지고(=KL 증가), 그만큼 reward 방향으로 더 이동한다(=reward 증가). \(\eta = 0.3\)처럼 작은 값을 고르면 KL 예산을 적게 쓰면서 SLERP가 만들어낸 방향의 일부만 취할 수 있다. 이 \(\eta\) 스윕이 바로 파레토 프런티어를 그리는 방법이고, WARP는 이 프런티어 위에서 \(\eta = 0.3\) 지점을 골라 다음 iteration의 새 초기화로 삼는다.

반복

이 세 단계를 \(I = 5\) iteration 동안 반복한다. iteration마다 \(M = 2\)개 정책을 학습하고, 학습 스텝 수는 9k → 7k → 5k로 점점 줄여간다(이전 iteration의 결과가 이미 좋은 초기화이므로 매번 처음부터 다시 많이 학습할 필요가 없다). 각 iteration이 끝나면 그 결과(\(\theta^\eta\))가 다음 iteration의 \(\theta_{init}\)이 되어, KL-reward 파레토 프런티어를 조금씩 밀어올린다.

Experiments

Gemma 7B 벤치마크

WARP로 정렬한 Gemma 7B는 Gemma 7B v1.1 대비 대부분의 벤치마크에서 개선을 보였다.

벤치마크 Gemma 7B v1.1 WARP
MBPP 39.0 45.4
GSM8K 55.6 66.8
MATH 25.6 31.0
MMLU 56.4 57.6
HumanEval 46.9 50.0

GSM8K에서 11.2점 오른 것이 가장 두드러지고, 코드·수학처럼 추론이 필요한 벤치마크 전반에서 개선폭이 크다. 반대로 이미 강했던 MMLU는 1.2점으로 개선폭이 작은데, 원래 여유가 적은 벤치마크일수록 weight averaging의 이득도 줄어드는 것으로 보인다.

선호도 승률

사람 선호도 기준으로도 WARP는 강한 baseline 대비 우위를 보였다.

비교 대상 선호 점수
Mistral 7B v2 +0.23
Mixtral 8x7B +0.26

EMA 앵커의 효과

\(\beta = 0.1\)의 EMA 앵커는 고정 앵커를 쓰는 두 변형(\(\beta = 0.01\), \(\beta = 0.001\))보다 여러 KL 수준에 걸쳐 일관되게 높은 reward를 기록했다. 이는 Stage 1에서 설명한 암묵적 annealing 효과가 실제로 작동한다는 근거다 — 동적 앵커 하나가 고정 앵커를 여러 개 튜닝해서 얻는 효과를 대신한다.

Conclusion

WARP는 KL 정규화와 reward 최적화 사이의 줄다리기를 \(\beta\) 하나로 풀려 하지 않는다. 대신 EMA 앵커로 정규화 자체를 동적으로 만들고, SLERP로 여러 정책의 task vector를 norm 손실 없이 합치고, LITI로 그 결과를 다시 초기화 쪽으로 당겨 파레토 프런티어 위의 원하는 지점을 고른다. 이 세 단계를 반복하면서 프런티어 자체를 iteration마다 밀어올린다는 것이 핵심 주장이다.

비용은 분명하다. iteration마다 여러 개의 RL 학습을 병렬로 돌려야 한다. 논문은 이를 “a feature rather than a bug”로 규정한다 — 사전학습이 compute를 늘려 성능을 얻듯, 정렬도 병렬 RL 학습을 늘려 compute를 스케일할 수 있다는 것이다.

#13 WARM이 reward 설계 쪽에서, WARP가 정책 학습 쪽에서 같은 weight averaging 아이디어를 각자의 방식으로 적용했듯, #28 BOND를 포함해 이 시기 DeepMind의 정렬 연구들은 서로 다른 컴포넌트에서 비슷한 직관을 공유한다. 실제로 #43 프론티어 모델의 reward 설계에서 다루듯, Gemma의 post-training 레시피는 BOND, WARM, WARP를 함께 조합한 결과물이다.


RLHF Reward 설계 시리즈

이 글은 RLHF Reward 설계 시리즈의 스물아홉 번째 글이다.

1부. 지형도

  1. Deep RL from Human Preferences (Christiano 2017) — 선호로 보상을 배우는 원형
  2. InstructGPT (Ouyang 2022) — RLHF 3단계 표준 레시피
  3. HH-RLHF (Bai 2022) — helpful·harmless preference model

2부. 스칼라 RM 해부

  1. Rethinking Bradley-Terry (2024) — reward 변환의 수학적 기반
  2. Secrets of RLHF II (2024) — 선호 데이터 노이즈와 RM 일반화
  3. Skywork-Reward (2024) — 데이터 큐레이션이 아키텍처를 이긴다
  4. ArmoRM (2024) — 다목적 분해와 MoE 게이팅
  5. Llama 2 (2023) — helpfulness·safety RM 분리 프로덕션 레시피
  6. RewardBench 2 (2025) — RM을 어떻게 평가할 것인가

3부. Reward Hacking

  1. Overoptimization Scaling Laws (2022) — Goodhart의 법칙 정량화
  2. Length Correlations in RLHF (2023) — 성능 향상의 얼마가 길이인가
  3. ODIN (2024) — 길이를 reward에서 분리
  4. WARM (2024) — weight averaging으로 hacking 방어

4부. 안전성 정렬

  1. Safe RLHF (2023) — 안전성을 reward가 아니라 제약으로
  2. Rule-Based Rewards (2024) — 안전 규칙을 reward로 직접 번역
  3. Deliberative Alignment (2024) — 안전 명세를 모델의 추론 안으로
  4. Shallow Safety Alignment (2024) — 정렬은 첫 몇 토큰에만 얹혀 있다
  5. OR-Bench (2024) — 과잉 거절을 어떻게 측정할 것인가

5부. reward를 정책으로

  1. PPO (2017) — clipped surrogate objective
  2. Secrets of RLHF I (2023) — PPO 학습 안정화 트릭
  3. GRPO / DeepSeekMath (2024) — value network를 버리다
  4. RLOO (2024) — REINFORCE로 충분한가
  5. DPO (2023) — reward를 없애면 어떻게 되는가
  6. SimPO (2024) — reference-free + 길이 정규화
  7. KTO (2024) — 선호 쌍 없이 이진 신호만으로
  8. GSPO (2025) — importance ratio를 시퀀스 단위로
  9. DAPO (2025) — 신호 없는 프롬프트를 버린다
  10. BOND (2024) — Best-of-N을 추론 비용 없이
  11. (현재 글) WARP (2024) — 정책을 weight space에서 병합

6부. Process & Verifiable Reward

  1. Let's Verify Step by Step (2023) — 과정 감독이 결과 감독을 이긴다
  2. Math-Shepherd (2023) — 사람 라벨 없는 PRM
  3. DeepSeek-R1 (2025) — RLVR, 규칙이 reward가 될 때

7부. Generative Reward Model

  1. Prometheus 2 (2024) — 오픈 평가자 모델과 rubric 조건부 평가
  2. Generative Verifiers (2024) — reward를 next-token prediction으로
  3. Generative Reward Models (2024) — GenRM과 선호 학습의 결합
  4. Self-Taught Evaluators (2024) — 사람 라벨 없이 judge를 키우다
  5. DeepSeek-GRM / SPCT (2025) — inference-time scaling

8부. 생각하는 Judge, 그리고 그 신뢰

  1. ReasonGRM (2025) — reasoning 능력을 judge에 이식
  2. J1 (2025) — RL로 judge를 생각하게 만들기
  3. Rubrics as Rewards (2025) — 비검증 도메인으로
  4. CriticEval (2024) — judge 자체를 어떻게 평가하나
  5. One Token to Fool LLM-as-a-Judge (2025) — GenRM도 뚫린다

9부. 실전 종합

  1. 프론티어 모델의 reward 설계 (2025~2026) — 열한 개 모델이 실제로 택한 것
  2. reward를 어떻게 설계할 것인가 — 시리즈를 관통한 RM 설계 원칙 한 장

본 시리즈는 44편으로 구성된다.

참고 문헌




Enjoy Reading This Article?

Here are some more articles you might like to read next:

  • reward를 어떻게 설계할 것인가 — RM 설계 실무
  • 프론티어 모델은 reward를 어떻게 설계했나
  • One Token to Fool: GenRM도 결국 뚫린다
  • CriticEval: judge를 채점하는 벤치마크
  • Rubrics as Rewards: 정답이 없는 도메인에 reward를 만드는 법