단계별 어드밴티지 선택을 통한 효율적인 추론의 안정화

LLM의 추론 능력을 높이기 위해 생성 길이를 늘리는 시도가 많지만, 이는 비용과 속도 면에서 비효율적이다. 이 논문은 추가적인 보상 모델 없이도 추론 과정을 압축하면서 학습 안정성을 확보하는 방법을 제시하여 효율적인 AI 추론 연구의 새로운 방향을 제시한다.

기존의 GRPO와 같은 강화학습 방식은 전체 추론 결과가 맞으면 그 안의 모든 단계를 긍정적으로, 틀리면 모든 단계를 부정적으로 평가하는 Rollout-level 방식을 사용한다. 하지만 정답 안에도 중복되거나 무의미한 단계가 섞여 있을 수 있고, 오답이라 하더라도 중간 과정은 논리적으로 완벽하지만 마지막에 컨텍스트 제한으로 잘린 경우가 존재한다. 이러한 '노이즈' 섞인 보상 신호는 학습을 불안정하게 만든다.

SAS는 이를 해결하기 위해 Gradient Descent 과정에서 각 추론 단계의 신뢰도(Confidence)를 계산한다. 정답 경로에서는 신뢰도가 낮은 단계를 '운 좋게 포함된 무의미한 단계'로 보고 학습 가중치를 0으로 설정해 무시한다. 반대로 오답 경로에서는 신뢰도가 높은 단계를 '잘렸지만 유효한 논리'로 판단하여 부정적인 보상으로부터 보호한다.

결과적으로 모델은 불필요한 군더더기 없이 핵심적인 추론 단계에만 집중하게 된다. 이는 마치 시험 공부를 할 때 단순히 정답 여부만 보는 것이 아니라, 내가 확실히 알고 푼 단계와 찍어서 맞춘 단계를 구분하여 학습 효율을 극대화하는 것과 같은 원리다.

전체적인 접근 방식은 GRPO(Group Relative Policy Optimization) 프레임워크를 기반으로 하되, 보상 할당 방식을 단계별(Step-level)로 세분화한다. 추론 과정은 이중 줄바꿈(

)을 기준으로 개별 단계로 분할되며, 각 단계에 대해 모델의 토큰 로그 확률 평균값을 사용하여 신뢰도 점수 c_j를 산출한다.

정답 롤아웃(Reward=1)의 경우, 신뢰도 점수가 낮은 하위 r 비율의 단계들에 대해 어드밴티지를 0으로 설정한다. [단계별 로그 확률 평균 계산 → 오름차순 정렬 → 하위 r 비율 선택 → Advantage=0 할당] 과정을 거쳐 불필요한 반복이나 확신 없는 단계를 강화하지 않도록 차단한다.

오답 롤아웃(Reward=0)의 경우, 신뢰도 점수가 높은 상위 r 비율의 단계들에 대해 어드밴티지를 0으로 설정한다. [단계별 로그 확률 평균 계산 → 내림차순 정렬 → 상위 r 비율 선택 → Advantage=0 할당] 과정을 통해, 컨텍스트 제한으로 인해 잘렸을 뿐 논리적으로는 우수한 중간 단계들이 부정적인 피드백을 받아 억제되는 것을 방지한다.

DeepScaleR-1.5B-Preview 모델을 베이스로 실험한 결과, SAS는 수학 벤치마크(AIME, MATH, AMC 등)에서 평균 Pass@1 54.54%를 기록하여 베이스라인인 GRPO-4K(53.61%)보다 높은 정확도를 보였다. 동시에 평균 토큰 수는 3,775개에서 3,407개로 약 10% 감소했다.

일반 추론 벤치마크(GPQA, LSAT, MMLU)에서도 SAS는 평균 38.30%의 정확도를 기록하며 베이스 모델(37.44%)을 앞질렀다. 특히 기존의 길이 제약 기법인 LAPO-I나 ThinkPrune-4k와 비교했을 때, 정확도와 효율성을 동시에 고려한 AES(Accuracy-Efficiency Score) 지표에서 0.46점을 기록하여 가장 우수한 성능을 입증했다.

Ablation Study를 통해 정답 경로의 마스킹과 오답 경로의 보호 메커니즘이 모두 유효함을 확인했다. 특히 오답 경로에서 유효한 단계를 보호하지 않을 경우 학습 초기부터 탐색 능력이 급격히 저하되는 현상이 관찰되었다.

SAS는 GRPO의 그룹 상대적 정규화(Group-relative normalization) 특성을 활용한다. 어드밴티지를 0으로 설정하는 행위는 정답 그룹 내에서는 다른 긍정적 단계들보다 낮은 가치를 부여하는 '억제' 효과를 내고, 오답 그룹 내에서는 다른 부정적 단계들보다 높은 가치를 부여하는 '보호' 효과를 내는 비대칭적 이득을 제공한다.

학습 안정성 측면에서 SAS는 Policy Entropy를 높게 유지하는 데 기여한다. 순수 GRPO-4K 학습 시 엔트로피가 급격히 붕괴하며 특정 패턴에 고착되는 현상이 발생하지만, SAS는 노이즈 섞인 신호를 필터링함으로써 모델이 더 다양한 추론 경로를 탐색할 수 있도록 돕는다.

구현 측면에서 SAS는 모델 아키텍처를 수정하지 않으며, 학습 시 약 17%의 시간 오버헤드만 발생시킨다. 이는 별도의 외부 보상 모델(PRM)을 호출하거나 다중 추론 경로를 생성해야 하는 다른 기법들에 비해 매우 가벼운 수준이다.

본 연구는 단일 베이스 모델(DeepScaleR-1.5B)을 대상으로 실험되었으며, 모델 크기나 사전 학습 방식이 다른 경우에도 동일한 경향성이 나타나는지에 대한 추가 검증이 필요하다. 또한 고정된 컨텍스트 길이(4K) 환경에서의 동작에 집중했으므로 다양한 컨텍스트 길이 변화에 따른 SAS의 거동은 향후 과제로 남겨져 있다.