세 LLM에서 GRPO 후 성능 저하와 KV cache 검증

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

저자는 동일한 synthetic arithmetic curriculum과 후학습 절차로 세 LLM을 비교했지만, GRPO 뒤 WikiText perplexity와 downstream 성능이 V1에서는 거의 유지되고 V2에서는 52%, V3에서는 5% 악화됐다고 보고했습니다. GRPO는 학습한 curriculum의 정답 생성은 익혔지만 GSM8K transfer에는 실패했고, 종료 보상 부재와 SFT·GRPO prompt format 차이가 결과를 교란했을 가능성이 있습니다. 모델별 변경점이 한꺼번에 바뀌고 이전 curriculum stage 재평가와 ablation이 없어 원인은 확정되지 않았습니다. 별도로 직접 구현한 GQA-aware KV cache는 logits 최대 차이 1.4e-06을 기록했고 100-token 생성 속도를 prompt 길이에 따라 3.7x에서 10.1x까지 높였습니다.

실용적 조언

GRPO 전후에 모든 curriculum stage를 다시 평가하면 일반 능력 저하와 sequential curriculum forgetting을 구분할 수 있습니다. 각 stage의 정답률과 생성 길이를 함께 기록해야 학습 과제 습득과 출력 제어 문제를 분리할 수 있습니다. 글의 현재 결과는 이 재평가가 빠져 있어 후속 실험의 우선순위가 높습니다.
Reward function에 정답 여부뿐 아니라 적절한 종료를 반영하고, SFT와 GRPO에서 동일한 prompt format을 유지하는 편이 좋습니다. 현재 reward는 parseable number가 출력되면 길이에 관계없이 보상하고, GRPO는 bare solver template를 사용했습니다. 이 두 조건을 맞추면 긴 생성과 format mismatch가 만드는 평가 교란을 줄일 수 있습니다.
V2와 V3를 비교할 때는 parameter count, token count, data mix, attention mechanism을 한 번에 하나씩만 바꾸는 ablation이 필요합니다. KL coefficient도 0.02 외의 값을 별도로 시험해 reference policy에서 벗어나는 정도와 perplexity 변화를 함께 측정해야 합니다. 전체 비용 제약이 있다면 작은 모델과 짧은 curriculum에서 먼저 조건별 경향을 확인하는 방식이 적합합니다.

섹션별 상세

저자는 raw PyTorch로 세 LLM을 처음부터 학습한 뒤 동일한 synthetic arithmetic curriculum, reward function, hyperparameters, KL coefficient를 적용해 SFT와 GRPO를 순차적으로 수행했습니다. V1은 353M parameters와 10B tokens, V2는 316M parameters와 10B tokens, V3는 672M parameters와 30B tokens로 구성됐으며 attention과 데이터 혼합도 모델마다 달랐습니다. Pre-training validation loss는 V1 2.8659, V2 2.7844, V3 2.5885로 낮아졌지만, 이 차이는 이후 GRPO 결과를 단순한 모델 크기 효과로 귀속하기 어렵게 만듭니다.

저자의 핵심 결과는 GRPO가 SFT 이후 WikiText word perplexity를 모든 모델에서 악화시켰다는 점입니다. V1은 51.31에서 51.40으로 거의 변하지 않았고, V2는 46.81에서 71.06으로 52% 상승했으며, V3는 32.11에서 33.65로 5% 상승했습니다. Downstream task도 같은 방향으로 움직였고 V3의 arc_easy 점수는 SFT에서 GRPO로 넘어가며 약 6점 낮아졌습니다.

GRPO는 학습에 사용한 산술 과제 자체는 습득시켰지만 GSM8K로의 transfer는 거의 일어나지 않았습니다. V3는 다섯 curriculum stage 중 네 개를 익혔고 다른 두 모델은 세 개씩 익혔지만, reward가 정답 숫자가 어디엔가 나타나는지만 확인하고 종료를 보상하지 않아 긴 풀이를 계속 생성하는 문제가 생겼습니다. 또한 GRPO에는 bare solver template를, SFT에는 chat format을 사용했으므로 일부 downstream 저하는 각 모델의 training distribution과 평가 형식이 어긋난 결과일 수 있습니다.

저자는 GRPO의 악화가 일반 능력 저하인지 sequential curriculum training에 따른 이전 단계 망각인지 아직 구분하지 못한다고 밝혔습니다. 이전 curriculum stage를 모델이 다음 단계로 진행한 뒤 다시 평가하지 않았기 때문에, 단계별 재평가와 KL coefficient 변경 실험이 원인 분리에 필요합니다. 세 모델 사이에서 parameter count, token count, data mix, attention mechanism을 동시에 바꿨고 전체 비용이 약 $750이어서 ablation을 수행하지 못했다는 점도 결과 해석의 제한입니다.

추론 측면에서는 GQA-aware KV cache를 요청별 cache object로 구현하고 동일한 고정 sequence를 full forward와 prefill-then-decode 방식으로 각각 계산했습니다. 두 방식의 logits 최대 차이는 1.4e-06으로 1e-4 tolerance보다 작았고, 100 tokens 생성 속도는 32-token prompt에서 3.7x, 128-token prompt에서 6.2x, 512-token prompt에서 10.1x 향상됐습니다. 따라서 학습 실험의 GRPO 결과와 별개로, 직접 작성한 KV cache의 수치 일치와 긴 prompt에서의 추론 속도 개선은 재현 가능한 구현 결과로 남았습니다.

용어 해설

그룹 상대 정책 최적화(GRPO): — GRPO는 여러 샘플의 보상 차이를 상대적인 기준으로 삼아 언어 모델의 정책을 업데이트하는 강화학습 방식입니다. 이 글에서는 SFT 모델을 기준 정책으로 두고 동일한 KL coefficient와 reward function으로 산술 커리큘럼을 학습했습니다.
지도 미세 조정(SFT): — SFT는 입력과 정답으로 구성된 예시를 사용해 모델의 다음 토큰 예측 능력을 직접 조정하는 후학습 단계입니다. 이 실험에서는 SFT 뒤에 GRPO를 적용하고 두 단계의 평가 결과를 비교했습니다.
키-값 캐시(KV cache): — KV cache는 긴 텍스트를 생성할 때 이미 계산한 Attention의 key와 value를 저장해 이전 토큰을 반복 계산하지 않게 하는 추론 최적화입니다. 글에서는 GQA를 고려한 요청별 캐시 객체를 직접 구현하고 전체 forward와 prefill-then-decode 결과를 대조했습니다.
그룹화된 쿼리 어텐션(GQA): — GQA는 여러 query head가 더 적은 수의 key·value head를 공유하도록 구성해 Attention의 메모리 사용량을 줄이는 방식입니다. V2와 V3는 GQA 4:1을 사용했고, KV cache 구현도 이 구조를 반영했습니다.
퍼플렉서티(perplexity): — perplexity는 언어 모델이 평가 텍스트의 다음 토큰을 얼마나 잘 예측하는지 나타내는 지표로, 값이 낮을수록 예측 불확실성이 낮습니다. 이 글에서는 WikiText word perplexity가 GRPO 뒤 V2에서 46.81에서 71.06으로 상승했습니다.
KL 발산 계수(KL coefficient): — KL coefficient는 강화학습 중인 정책이 기준 정책에서 지나치게 멀어지지 않도록 두 정책의 분포 차이에 패널티를 부여하는 가중치입니다. 세 모델 모두 SFT 정책을 고정된 reference로 사용했고 값은 0.02였습니다.

언급된 도구

PyTorch중립

raw PyTorch로 LLM을 처음부터 학습하고 KV cache를 직접 구현하는 데 사용했습니다.

lm-evaluation-harness중립

동일한 task version과 shot count로 WikiText word perplexity 및 downstream task를 평가하는 데 사용했습니다.

언급된 리소스

문서Weights

DemoPlayground

GitHubCode

문서full writeup