Soft Adaptive Policy Optimization 논문 리뷰

Qwen에서 낸 SAPO 논문이다. 이 논문을 이해하기 위한 몇 가지 preliminaries를 확인해보자.

Preliminaries

Proximal Policy Optimization &  GRPO

PPO를 설명하기 위해서는, 왜 PPO까지 가게 되었는지에 대해서 이해해야 한다. 우선 policy update의 과정에서 data reuse를 하기 위해서는 Importance sampling이 필요하다. 이러면 어떤 문제점이 생기는지 아래의 그림을 보면서 이해해보자.

(옛날에 발표할 때 쓴 ppt중 일부임)

따라서 trajectory 단위로 importance sampling을 하는 것은 사실상 불가능에 가깝기 때문에, step 단위로 importance sampling을 하는 대안이 생겨난다.

(이거도 발표용...)

이렇게 되는 경우 policy의 stationary distribution이 업데이트 전과 후에 거의 달라지지 않는다는 가정이 필요하게 된다. 따라서 업데이트할 때 policy update의 크기를 제한해야 하며, 결국 TRPO에서 natural policy gradient를 통해 이론적으로는 문제를 해결하게 된다.

 

그러나 TRPO의 방법은 실제로는 큰 수치적 불안정을 일으키며 구현도 사실상 까다롭기 때문에 단순하면서 강력한 대안이 요구된 것이고, 이러한 수요에 따라 나온 것이 Proximal Policy Optimization이다.

엄청 복잡한 아키텍처 없이, PPO는 클리핑만으로도 좋은 성능을 가질 수 있음을 보였고 강화학습의 대표적인 알고리즘 중 하나로 자리잡게 된 것이다. 이에 나아가 PPO는 InstructGPT부터 시작하여 LLM을 강화학습으로 학습하는 기초가 되었다.

 

PPO를 응용한 LLM+RL 알고리즘 중 대표적인 것은 GRPO다.

GRPO는 하나의 query에 대한 여러 응답을 바탕으로 policy gradient를 한다는 점에서 분산을 낮춘다는 장점도 있지만, 동시에 GAE같은 전통적인 advantage estimation을 없애고 group 응답의 평균보다 얼마나 좋은지로 advantage를 구성한다는 점에서 instability를 줄인다는 장점이 있다.

 

MoE with RL

MoE로 PPO를 하면 이야기가 완전히 달라진다. 표면적으로는 token 당 policy의 ratio는 $r_t = \frac{\pi_\theta (a_t | s_t)}{\pi_{\text{ref}} (a_t | s_t)} $가 되어, 지금 현재의 토큰만 조금 조정한다는 의미를 가지지만 MoE에서는 단순히 조정으로 끝나지 않는다.

 

MoE LLM에서 한 토큰의 확률은 사실:

$$ \pi(a_t | s_t) = \sum_{e\in\mathcal{E}} \pi(e | s_t) \pi(a_t | s_t, e) $$

로, 어떤 expert를 사용하고, 그 expert가 무엇을 말할지 전부 업데이트 한다. 단일 LLM이었을 때에는 negative update가 그나마 출력에 사용된 토큰의 logit을 낮추고, 나머지를 전부 높이는 문제밖에 없었지만 MoE의 경우 다른 expert의 logit마저도 올리게 되어 routing probability까지 망가뜨리게 된다는 문제점이 있다.

 

특히, routing probability는 discrete하며 non-smooth, 즉 일정 수치가 넘으면 활성화되는 구조이기 때문에 작은 업데이트만으로도 큰 점프가 일어날 가능성이 높다. 또한, 토큰 $t$에서 expert가 바뀌면, hidden state가 바뀌고, 그에 따라서 이후의 state가 크게 달라지게 된다. 이는 앞서 알아보았던 stationary distribution ratio을 1로 두는 가정과 즉시 위반되는 결과까지 이어지는 것이다.

 

결국 MoE LLM에서 RL을 사용하기 위해서는 토큰 단위가 아닌 시퀀스 또는 청크 단위로 RL을 적용해야 하는 과제가 새로 생기게 되는 결과로 이어진다.

 

SAPO

SAPO에서 시퀀스 단위의 정책 변화를 고려한다는 것은, 전통적인 per-trajectory importance sampling과 PPO에서의 per-token importance sampling 사이의 절충 지점으로 이해할 수 있다.

 

전체 trajectory에 importance weight를 곱하는 방식은 분산 폭발이나 소실을 초래하기 쉬워 실제로 사용하기 어렵다. 반면, 토큰 단위 importance weight는 안정적이지만, prefix 수준에서 누적되는 state distribution mismatch를 충분히 반영하지 못한다.

 

SAPO는 이러한 두 극단 사이에서, 시퀀스 규모에서 정책 변화의 누적 효과를 평가하면서도 hard trajectory importance sampling을 명시적으로 수행하지 않도록 설계된 soft gating 메커니즘을 사용한다. 이를 통해 정책 변화가 지나치게 큰 경로는 감쇠하면서도, 비교적 on-policy에 가까운 부분에서 발생하는 학습 신호는 유지할 수 있다.

 

물론 SAPO는 MoE LLM에서 강화학습의 근본적인 구조적 문제를 해결하는 방법은 아니다. 다만 MoE 구조로 인해 학습 과정에서 발생하는 불안정한 업데이트가 치명적인 붕괴로 이어지는 것을 완화하는 안정화 기법으로 이해할 수 있다.

 

SAPO의 목적식은 다음과 같다:

여기서. $r_{i,t}(\theta)=\frac{\pi_\theta(a_t | s_t)}{\pi_{\text{ref}}(a_t | s)_t)} $이다. 핵심은 $f$이다. 만약 $r$이 지나치게 크거나 작다면 새로운 policy가 해당 action을 지나치게 밀어버리고 있으니, sigmoid로 중심에 가깝게 다가오게 해주는 것이다. Gradient 전파 과정을 보면 더욱 쉽게 이해할 수 있는데, sigmoid의 gradient는 중심에서 최대가 되고 중심에서 멀어질수록 gradient가 감쇠하게 되어 그 전파 신호가 줄어들게 되는 것이다.

즉 기존과 같이 sequence level로 보상을 일관적으로 제공하되, 문제가 되는 토큰만 선택적으로 감쇠하는 trust region 기법인 것이다.

 

그런데 의문은, 그러면 token level로 clipping을 적용한 PPO, GRPO에 비해 어떤 대단한 이득이 있냐는 것이다. SAPO가 전체 trajectory에 대한 보상을 통해 업데이트를 하다가도 특정 이상 상황에서는 chunk level로 업데이트를 할 수 있다는 장점이 있기 때문이다. 이렇게 말을 하면... 사실 내가 든 의문은 PPO나 GRPO도 clipping을 통해서 차단하기 때문에 비슷한 것 아니냐는 생각이 들었었는데, 그럼에도 정보 자체를 제거하는 것보다는 더 합리적인 방법이며 reward 정합성에도 맞기 때문이라고 한다.

특히 수학적으로, SAPO가 일정 조건 하에서 sequence chunk level rewarding과도 이어질 수 있음을 논문에서 보였으니 확인해보면 좋을 것 같다.

 

논문에서 Positive 응답과 Negative 응답에 대한 서로 다른 temperature을 부여하는지에 대한 이유가 위의 글에서 드러난다.

만약 긍정적 응답의 경우에는, 해당 토큰들의 logit이 높아지고 다른 토큰들의 logit이 낮아지는데, 이미 해당의 경우에서 다른 token들의 logit이 이미 작기 때문에 낮아지는 것에 대한 패널티가 크게 존재하지 않는다.

그러나 negative 응답에 대해서는 다른 토큰들의 logit을 높여주기 때문에, 업데이트 자체가 "억제"라는 기본 목적을 달성하지 못할 수 있게 되는 것이다. 따라서 negative의 temperature을 positive보다 높게 설정한다.

 

논문에서 아래의 실험을 통해서 그 결과를 경험적으로도 보여준다:

이번 논문 리뷰는 여기서 마치도록 하겠다.