Online Causal Kalman Filtering for Stable and Effective Policy Optimization

이 논문은 토큰 수준의 중요도 샘플링 비율의 높은 분산으로 인한 정책 최적화의 불안정성을 해결하기 위해, 과거 토큰의 상태에 기반하여 오프라인으로 추정된 중요도 비율을 칼만 필터링을 통해 실시간으로 보정하는 '온라인 인과적 칼만 필터링 기반 정책 최적화 (KPO)' 방법을 제안하고 수학 추론 작업에서 뛰어난 성능을 입증합니다.

핵심

🎓 비유: "혼란스러운 교실과 현명한 담임 선생님"

1. 문제 상황: 소음에 휩싸인 교실

인공지능 (LLM) 이 수학 문제를 풀면서 스스로 학습할 때, 마치 수천 명의 학생이 동시에 떠드는 교실과 같습니다.

• 기존 방식 (GRPO 등): 선생님이 "이 학생이 맞았으니 점수를 주고, 틀렸으니 깎아라"라고 할 때, 학생 한 명 한 명의 반응 (각 단어) 을 개별적으로 보다가, 때로는 너무 큰 소음 (변동성) 때문에 혼란이 생깁니다.
  • 예: "이 단어는 확실히 맞는데, 바로 다음 단어는 실수해서 점수가 급락했다"라고 하면, AI 는 "내가 뭘 해야 할지 모르겠다"며 학습이 불안정해집니다. (학습이 붕괴됨)
• 기존 해결책 (시퀀스 레벨): "아, 이 학생은 전체적으로 잘했으니 전체 점수만 평균내자!"라고 합니다. 하지만 이렇게 하면 세부적인 맥락 (어디서 잘하고 어디서 못했는지) 을 무시하게 되어, 정교한 학습이 어렵습니다.

2. 새로운 발견: "갑작스러운 기분 전환"은 이상하다

저자들은 데이터를 분석해보니, AI 가 학습할 때 단어 하나하나의 점수 변화가 너무 자주, 너무 급격하게 바뀌는 것을 발견했습니다.

• 마치 학생이 "지금 1 초는 천사, 1 초 뒤는 악마, 또 1 초 뒤는 천사"처럼 기분 전환을 너무 자주 하는 것과 같습니다.
• 실제로는 AI 가 어떤 논리 흐름 (추론 과정) 을 따라갈 때는 그 흐름이 일정하게 유지되어야 하는데, 계산상의 오류 (소음) 때문에 신호가 요동치는 것입니다.

3. 해결책: KPO (칼만 필터링을 쓴 현명한 담임)

이 논문은 **'온라인 인과 칼만 필터링 (KPO)'**이라는 새로운 방법을 제안합니다. 이를 현명한 담임 선생님에 비유해 볼까요?

• 과거의 경험과 현재를 연결: 담임 선생님은 학생이 "지금 실수했다"고 해도, "아까까지 잘하고 있었으니 이건 일시적인 실수일 뿐이야"라고 판단합니다.
• 소음 제거 (Noise Smoothing): 학생의 급격한 기분 전환 (소음) 은 무시하고, 전체적인 흐름 (구조) 을 따라가게 합니다.
  • 예: "이 학생은 10 단어를 연속해서 잘 풀고 있었어. 11 번째 단어에서 실수했더라도, 전체적인 흐름은 '잘하고 있는 중'이야."
• 미래는 모른다 (인과적): 담임 선생님은 과거와 현재의 정보만 가지고 미래를 예측합니다. (미래의 정답을 미리 알 수 없으므로) 그래서 AI 가 단어를 하나씩 생성해 나가는 방식과 완벽하게 맞습니다.

4. 결과: 더 안정적이고 똑똑한 학습

이 방법을 적용한 결과:

• 학습이 안정됨: AI 가 "내가 뭘 해야 할지 모르겠다"며 공황에 빠지는 일이 줄었습니다.
• 성능 향상: 특히 복잡한 수학 문제를 풀 때, 기존 방법들보다 훨씬 높은 점수를 받았습니다.
• 맥락 유지: 전체 점수만 평균내는 게 아니라, "어디서 잘하고 어디서 실수했는지"라는 세부적인 맥락도 살려냈습니다.

---

💡 핵심 요약 (한 줄 정리)

"AI 가 학습할 때 단어 하나하나의 점수 변동이 너무 심해 혼란스러워하는 문제를, '과거의 흐름을 기억하며 소음만 걸러내는 현명한 필터 (칼만 필터)'를 통해 해결함으로써, 더 안정적이고 똑똑한 학습을 가능하게 했다."

이 기술은 AI 가 더 복잡한 문제 (수학, 논리 등) 를 풀 때, 실수를 반복하지 않고 차분하게 추론할 수 있도록 도와주는 '학습의 안정장치' 역할을 합니다.

기술 요약

1. 문제 정의 (Problem Statement)

대규모 언어 모델 (LLM) 에 대한 강화학습 (RL) 은 사전 학습의 한계를 넘어 성능을 극대화하는 핵심 방법론이지만, 대규모 스케일에서 **정책 최적화 (Policy Optimization)**의 불안정성 문제가 발생합니다.

• 고분산 토큰 수준의 중요도 샘플링 (IS) 비율: 기존 방법 (GRPO 등) 은 토큰 단위로 중요도 샘플링 (Importance Sampling, IS) 비율을 계산합니다. 그러나 오프-폴리시 (off-policy) 설정에서 이 IS 비율은 토큰 간에 매우 높은 분산과 구조적 불일치를 보입니다.
• 구조적 불일치의 원인: 시퀀스 내의 인접 토큰들은 의미적으로 일관되어야 하지만, 실제 계산된 IS 비율은 짧은 구간에서도 급격하게 변동하거나 (Switching frequency 높음), 짧은 런 길이 (Run-length) 를 보입니다. 이는 인접 토큰 간의 정책 기울기 (Policy Gradient) 업데이트를 왜곡시키고, 결국 **학습 붕괴 (Training Collapse)**나 엔트로피 감소 (Entropy Collapse) 를 초래합니다.
• 기존 방법의 한계:
  • 시퀀스 수준 고정 비율 (GSPO, GMPO): 모든 토큰에 동일한 시퀀스 수준의 IS 비율을 적용하여 분산을 줄이지만, 시퀀스 내부의 국소적 구조 (Local Structure) 를 무시하여 세밀한 학습 신호를 잃습니다.
  • 개별 토큰 조정: 소프트 게이트나 비대칭 처리 등을 사용하지만, 토큰 간의 시간적 (Temporal) 구조를 고려하지 않아 노이즈를 완전히 제거하지 못합니다.

2. 제안 방법: KPO (Key Methodology)

이 논문은 **온라인 인과 칼만 필터링 (Online Causal Kalman Filtering)**을 기반으로 한 새로운 정책 최적화 방법인 KPO를 제안합니다.

• 핵심 아이디어: 토큰 단위의 IS 비율을 단순한 노이즈가 아닌, **시간에 따라 진화하는 잠재 상태 (Latent State)**로 모델링합니다. 칼만 필터를 적용하여 과거와 현재 토큰의 정보만을 사용하여 (인과적, Causal) IS 비율을 실시간으로 추정하고 평활화 (Smoothing) 합니다.
• 칼만 필터링 모델:
  • 상태 공간 모델 (State-Space Model):
    • 잠재 상태 ($\rho_t$): 이상적인 IS 비율 (로그 공간) 로, 인접 토큰 간에 부드럽게 변화한다고 가정합니다 ($\rho_t = \rho_{t-1} + \eta_t$).
    • 관측 모델 ($z_t$): 실제 계산된 IS 비율 ($\log r_t$) 로, 관측 노이즈 ($\epsilon_t$) 를 포함합니다.
  • 필터링 과정:
    1. 예측 (Prediction): 이전 상태의 평균과 분산을 기반으로 다음 토큰의 상태를 예측합니다.
    2. 칼만 이득 (Kalman Gain): 예측 불확실성과 관측 노이즈의 비율을 계산하여, 관측값과 예측값 중 어느 쪽에 더 가중치를 둘지 결정합니다.
    3. 업데이트 (Update): 관측된 IS 비율과 칼만 이득을 결합하여 현재 토큰의 필터링된 IS 비율 ($\hat{\rho}_t$) 을 추정합니다.
• 정책 최적화 목적 함수: 필터링된 IS 비율 ($e^{\hat{\rho}_t}$) 을 GRPO 와 유사한 목적 함수의 IS 비율 항에 대입합니다. 이를 통해 토큰 단위의 이질성 (Heterogeneity) 은 유지하되, 국소적 노이즈 스파이크는 제거하여 안정적인 업데이트를 유도합니다.

3. 주요 기여 (Key Contributions)

1. 실증적 발견 (Empirical Revelation): 토큰 수준의 오프-폴리시 편차가 **구조적 불일치 (Structural Inconsistency)**를 보인다는 것을 처음으로 규명했습니다. 즉, 인접 토큰 간에 IS 비율이 급격히 변동하여 학습을 불안정하게 만든다는 사실을 통계적 지표 (Switching frequency, Run-length 등) 로 증명했습니다.
2. 새로운 최적화 알고리즘 (KPO): 시퀀스 내 IS 비율에 인과적 칼만 필터를 적용하여 국소적 변동을 평활화하면서도 시퀀스 내부의 구조적 일관성을 보존하는 방법을 제안했습니다. 이는 기존 시퀀스 평균화 방식과 토큰 단위 노이즈 방식의 단점을 모두 보완합니다.
3. 강력한 실험 결과: 다양한 수학 추론 벤치마크에서 기존 최첨단 (SOTA) 방법들 (GRPO, GMPO, GSPO) 보다 우수한 성능과 학습 안정성을 입증했습니다.

4. 실험 결과 (Experimental Results)

• 벤치마크: AIME'24, AIME'25, AMC'23, MATH500, OlympiadBench 등 6 가지 어려운 수학 추론 데이터셋에서 평가.
• 성능:
  • KPO-Clipped는 대부분의 벤치마크에서 avg@16과 pass@16 지표에서 SOTA 를 기록했습니다. 특히 복잡한 추론이 필요한 AIME'24 와 AIME'25 에서 GSPO 대비 큰 성능 향상 (예: AIME'24 avg@16 32.70 → 37.91) 을 보였습니다.
  • KPO-Unclipped도 AIME'24 pass@16 에서 66.67 로 최상위 성능을 기록했습니다.
• 학습 안정성 (Training Dynamics):
  • 보상 (Reward): KPO 는 학습 전체에 걸쳐 지속적으로 보상을 향상시키는 반면, GRPO 는 200 스텝 이후 발산하고 GSPO/GMPO 는 초기에 수렴하여 정체되는 모습을 보였습니다.
  • 엔트로피 (Entropy): GRPO 는 초기에 엔트로피가 0 에 수렴하여 탐색이 멈추는 반면, KPO 는 높은 엔트로피를 유지하며 안정적인 탐색을 가능하게 했습니다.
  • 클립 비율 (Clip Fraction): KPO 는 시퀀스 전체를 클립하는 것이 아니라 필터링된 구조에 기반하여 필요한 토큰만 선택적으로 클립하여 더 효율적인 업데이트를 수행했습니다.
• 분포 분석: 칼만 필터링 적용 후 IS 비율의 **런 길이 (Run-length)**가 급격히 증가하고 (1.64 → 119.95), **전환 빈도 (Switch frequency)**가 급감 (0.43 → 0.01) 하여, 토큰 상태가 긴 구간 동안 일관되게 유지됨을 확인했습니다.

5. 의의 및 결론 (Significance)

• 학습 안정성 확보: 대규모 LLM 의 강화학습에서 발생하는 고분산 IS 비율 문제를 해결하여, 학습 붕괴 없이 안정적으로 확장 가능한 정책 최적화를 가능하게 합니다.
• 구조 인식 (Structure-Aware) 접근: 단순히 노이즈를 제거하는 것을 넘어, 시퀀스 내 토큰 간의 **국소적 일관성 (Local Coherence)**을 보존하는 새로운 관점을 제시했습니다. 이는 복잡한 추론 작업 (Chain-of-Thought) 에서 특히 중요한 요소입니다.
• 실용성: 경량화되어 기존 RL 파이프라인 (GRPO 등) 에 쉽게 통합 가능하며, MoE(Mixture-of-Experts) 아키텍처나 학습 - 추론 불일치가 발생하는 환경에서도 효과적입니다.

이 논문은 강화학습의 안정성을 높이기 위해 통계적 필터링 기법 (칼만 필터) 을 도입한 선구적인 연구로, 대규모 언어 모델의 고도화된 추론 능력 향상에 중요한 기여를 합니다.