LoGRA가 저랭크 그래디언트 스케치로 LLM 강화학습 메모리 장벽을 낮춘다
LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches
무엇인가
강화학습은 LLM 성능을 끌어올리는 핵심 후처리 단계가 됐지만, 메모리 요구량이 확산의 걸림돌로 남아 있다. 논문이 드는 예시는 명확하다. 7B Qwen 모델의 BF16 가중치는 약 14GB인데, Adam의 FP32 모멘트 버퍼 두 개가 56GB를 더 쓴다. 그래디언트와 활성화는 그전에 계산에 들어가지도 않는다. 즉 모델을 돌릴 수 있는 하드웨어가 그 모델을 학습시킬 수 있는 하드웨어는 아닐 수 있다는 것이 이 논문이 푸는 문제다.
어떻게 동작하나
LoGRA의 핵심은 전체 그래디언트 행렬을 만들지 않고 저랭크 스케치로 압축해 누적하는 것이다. 가중치 행렬 W∈R^{d×k}에 대해 랭크 r≪min(d,k)과 랜덤 투영 행렬 A∈R^{r×k}를 고르고, 그래디언트를 S=GA^T∈R^{d×r}로 표현한다. A는 배치와 무관하게 샘플링되어 그래디언트 누적 동안 고정되며 학습되지 않는다. 기본값은 ±1/√r 스케일의 독립 랜덤 부호(Rademacher)다. 구현상으로는 전체 그래디언트를 만든 뒤 압축하는 게 아니라, 마이크로배치마다 투영된 기여분을 fp32로 S에 직접 누적한다. 업데이트는 W←W−αηUA 형태로, S를 옵티마이저가 조정한 U와 A의 곱을 가중치에 바로 병합한다. 어댑터 분기를 따로 유지하지 않는다. 적용 대상은 어텐션과 MLP 투영 행렬이고 나머지 파라미터는 고정한다. A를 매 스텝 새로 뽑으면 연속 업데이트가 서로 다른 부분공간을 쓰게 되어 합의 랭크가 r 이하로 묶이지 않는다. 메모리 효과는 d=k=4096, r=64일 때 fp32 저장 기준 64MiB가 1MiB로 줄어드는 수준이다. 같은 압축 표현은 롤아웃 정책 동기화에도 재사용된다. 롤아웃 엔진이 시드로 A를 재생성하고 αηU만 받아 적용하므로, 전송량이 전체 d×k 업데이트 대신 O(dr) 값에 시드와 메타데이터를 더한 크기가 된다.
무엇과 다른가
압축만으로는 안정적인 업데이트가 보장되지 않는다는 것이 저자들의 판단이다. 파라미터 공간에서 작은 변화도 정책의 출력 확률은 크게 흔들 수 있기 때문이다. 그래서 예측 KL 스텝 제어를 결합한다. 현재 배치에서 일부 응답을 골라, 제안된 업데이트가 다음 토큰 확률을 얼마나 바꿀지 q(D)로 추정한다. 구체적으로는 각 컨텍스트에서 토큰 점수(logit)가 업데이트 방향 D를 따라 변하는 비율의 분산을 현재 확률로 가중해 평균낸다. 예측 KL은 스텝 크기의 제곱에 비례하므로, KL 예산 δ에 맞춰 α=min{α_max, √(δ/q(D))}로 업데이트를 재조정한다. 예컨대 q(D)=4δ인 제안은 원래 크기의 절반 이하로 줄어든다. 이는 TRPO의 신뢰영역 발상에 기대를 둔 것으로, 국소 근사이기 때문에 실제 KL은 추정치와 다를 수 있다고 논문은 명시한다.
어떻게 쓰나
실험은 Qwen2.5-Math-1.5B/7B와 Qwen3.8-27B를 수학 추론 과제에서 검증 가능한 보상의 PPO로 학습한다. 배치당 128 프롬프트, 응답 1개, 4,096 토큰 컨텍스트, 단일 노드 8×H100 80GB(4장은 FSDP 학습, 4장은 생성 엔진) 구성이다. DAPO-Math-7.5K에서 LoGRA는 평균 학습 메모리를 1.5B에서 9.18→7.18GiB(21.8%), 7B에서 31.82→17.29GiB(45.7%)로 줄였다. 시드 평균 피크 기준으로는 1.5B가 9.19→8.66GiB, 7B가 31.85→22.26GiB다. 성능은 1.5B에서 최고 Pass@1이 67.87%로 dense의 63.77%를 앞섰고 Pass@4도 80.33%에서 82.47%로 올랐다. 7B에서는 72.33% 대 72.48%(Pass@1), 85.13% 대 85.33%(Pass@4)로 사실상 동등했다. 27B는 LoGRA가 Pass@1 71.52%, Pass@4 81.87%에 도달한 반면 dense Adam은 첫 옵티마이저 상태 할당에서 GPU 메모리가 부족해 업데이트를 끝내지 못했다. 처리량은 양쪽이 학습을 완주한 구간에서 1.5B 55.44 대 54.28 updates/h, 7B 25.73 대 25.28 updates/h로 비슷했다.
전제와 한계
단일 8-GPU 노드에서의 장기 학습도 보고된다. Reasoning-Gym Hard 혼합(학습 프롬프트 38,948개, 8개 과제의 held-out 193문제)에서 랭크 256 스케치와 어닐링 KL 예산으로 Qwen3.8-27B를 학습시켜, held-out 매크로 점수가 39.69%에서 62.94%로 상승하고 1,060스텝에서 65.52%를 찍었다. 마지막 다섯 번의 평가는 59.11~64.16% 범위로 지속적 붕괴가 없었다. 같은 구간에서 평균 응답 길이는 8,312토큰에서 4,705토큰으로, 잘림률은 34.72%에서 12.95%로 줄었다.
절제 실험은 설계 선택의 무게를 보여준다. 랭크를 4에서 256으로 올리면 Pass@1이 3.02~3.95%포인트 개선됐고, 그래디언트 재구성의 프로베니우스 코사인 유사도는 랭크 4의 약 0.046에서 랭크 256의 0.348로 올랐다. 반면 기저를 매 업데이트 갱신하는지, 투영 분포가 Rademacher인지 Gaussian인지의 차이는 레이어-랭크 셀당 최대 0.0014, 0.0008 수준으로 미미했다. LoRA와의 비교에서는 LoGRA가 평균 7.18GiB, 피크 8.66GiB로 LoRA의 13.21GiB, 13.38GiB보다 메모리를 적게 썼지만, LoRA가 평균 Pass@1과 정책 학습 호출 시간에서 우세했고 LoGRA는 평균 Pass@4에서 0.40%포인트 앞섰다. 저자들은 이를 메모리가 제한 자원일 때의 대안 근거로 제시할 뿐, 정확도나 속도 우위를 입증한 것은 아니라고 선을 긋는다. KL 예산 절제에서는 700~800스텝 구간 평균 Pass@1이 제어 없음 70.30%, 큰 상수 예산 63.43%, 작은 상수 예산 71.03%, 어닐링 71.47%였다.
실무 관점에서 이 논문이 주는 신호는 두 가지다. 첫째, RL 후처리에서 메모리 병목의 상당 부분은 가중치가 아니라 그래디언트와 옵티마이저 상태이며, 이를 저랭크 표현으로 바꾸면 같은 GPU 예산으로 더 큰 모델을 학습 범위에 넣을 수 있다. 둘째, 압축은 그 자체로 학습을 불안정하게 만들 수 있어 KL 예산 같은 스텝 크기 규제와 함께 써야 한다는 것이다. 다만 보고된 메모리 수치는 생성 메모리를 제외한 학습 메모리 기준이고, 27B에서는 dense 베이스라인이 OOM으로 실패해 메모리 절감률이나 처리량 비교 자체가 측정되지 않았다. 도입을 검토한다면 자신의 랭크·KL 예산·FSDP 샤딩 구성에서 재현해 보는 것이 먼저다.
저자들이 밝힌 한계도 분명하다. LoGRA 단독으로는 안정적 업데이트가 보장되지 않으며, 예측 KL은 국소 근사라 실제 KL과 어긋날 수 있다. KL 예산 절제는 시드 2개 결과로, 잘 보정된 KL 제어가 이점을 준다는 시사는 있지만 안정성에 필수라거나 어닐링 자체가 이득의 원인이라고 단정하지 않는다. 기저 갱신이 이득을 준다는 근거도 짧은 학습 구간에서는 일관되게 나타나지 않았다. 또한 실험은 공개 모델과 벤치마크에 한정되며, 저널은 일반 목적 학습 기법과 마찬가지로 하위 활용에 적절한 안전장치와 평가가 필요하다고 덧붙인다.