Value Gradient Flow를 통한 강화학습

기존의 행동 규제 강화학습은 복잡한 정책 파라미터화나 보수적인 샘플링에 의존해 대규모 생성 모델로 확장하기 어려웠다. 이 논문은 최적 운송 이론을 도입해 가치 함수의 그래디언트만으로 정책을 유연하게 조정함으로써, 오프라인 RL과 LLM 정렬에서 SOTA 성능을 달성하고 테스트 시점의 적응형 스케일링까지 가능하게 한다.

강화학습에서 정책을 학습시킬 때 가장 큰 문제 중 하나는 모델이 학습 데이터 범위를 벗어난 영역에서 가치를 과대평가하여 잘못된 행동을 선택하는 것이다. 이를 막기 위해 기존에는 정답지 역할을 하는 참조 분포(Reference Distribution)에서 너무 멀어지지 않도록 KL Divergence 같은 페널티를 주었으나, 이는 가치 학습과 정책 개선 사이의 균형을 맞추기 어렵게 만들고 대규모 모델에서는 계산 비용이 매우 높았다.

VGF는 이 문제를 '입자들을 좋은 곳으로 조금씩 밀어주는 과정'으로 해석한다. 먼저 참조 분포에서 여러 개의 샘플(입자)을 뽑은 뒤, 가치 함수(Q-function)의 그래디언트를 따라 이 입자들을 가치가 더 높은 방향으로 조금씩 이동시킨다. 이 과정은 물리학의 흐름(Flow) 개념과 유사하며, 입자들이 이동하는 거리나 횟수 자체가 자연스럽게 행동 규제(Regularization) 역할을 수행하게 된다.

결과적으로 VGF는 정책을 정의하는 복잡한 수식 없이도 입자들의 이동만으로 최적의 행동을 찾아낸다. 이는 마치 지형도에서 경사가 높은 곳을 향해 공을 굴리는 것과 같으며, 이동 횟수를 조절함으로써 참조 분포를 얼마나 따를지 혹은 새로운 고가치 영역을 얼마나 탐색할지를 자유롭게 결정할 수 있게 한다.

VGF는 행동 규제 RL을 참조 분포 $\mu$에서 볼츠만 분포 형태의 최적 정책 $\pi^*_R \propto \exp(R/\alpha)$로 질량을 옮기는 최적 운송(Optimal Transport) 문제로 공식화한다. 이를 해결하기 위해 Wasserstein 메트릭 공간에서의 그래디언트 흐름을 활용하며, 연속적인 시간 진화를 이산적인 JKO(Jordan-Kinderlehrer-Otto) 스킴으로 근사하여 입자 업데이트 규칙을 도출한다.

핵심 메커니즘은 입자 기반의 SVGD(Stein Variational Gradient Descent)와 유사한 업데이트 식을 사용한다. $l$번째 단계의 입자 $a_i^{(l)}$는 가치 함수의 그래디언트 $abla_a R(s, a)$와 입자 간의 반발력을 제공하는 커널 함수 $k(a_j, a)$의 조합에 의해 $a_i^{(l+1)}$로 갱신된다. [현재 입자 위치와 가치 그래디언트를 입력으로] → [커널로 평활화된 가치 상승 방향과 입자 간 거리를 유지하는 척력을 계산하여] → [입자의 다음 위치를 결정하고] → [이 과정이 반복될수록 입자들은 참조 분포 근처의 고가치 영역으로 수렴한다].

LLM 환경에서는 토큰이 이산적이므로 연속적인 임베딩 공간이나 잠재 공간(Latent Space)에서 VGF를 수행한 뒤 마지막에 디코딩하는 방식을 취한다. 보상 모델의 그래디언트를 체인 룰(Chain Rule)을 통해 입력 임베딩까지 역전파하여 텍스트 생성 과정을 제어하며, 이는 PPO와 같은 고비용의 강화학습 최적화 없이도 효과적인 정렬을 가능하게 한다.

D4RL 벤치마크에서 VGF는 MuJoCo 로코모션 및 AntMaze 작업 전반에 걸쳐 Gaussian 정책이나 Diffusion 기반 모델들을 능가했다. 특히 AntMaze-Large와 같은 어려운 작업에서 기존 방법론들이 50점 미만의 점수를 기록할 때 VGF는 80점 이상의 높은 성공률을 보였다. OGBench 실험에서도 Humanoid Maze와 같은 복잡한 고차원 제어 작업에서 타 모델 대비 월등한 성능과 빠른 수렴 속도를 입증했다.

RLHF 작업인 TL;DR 요약 및 Anthropic-HH 데이터셋에서도 VGF의 우수성이 확인됐다. GPT-4를 이용한 승률 평가에서 VGF는 PPO(57.3%)나 DPO(61.2%)보다 높은 68.1%의 승률을 기록했다. 이는 VGF가 보상 모델의 정보를 직접적인 그래디언트 흐름으로 활용함으로써 보상 해킹(Reward Hacking)을 억제하면서도 더 나은 품질의 텍스트를 생성함을 의미한다.

Ablation Study를 통해 학습 시의 흐름 단계($L_{train}$)와 추론 시의 단계($L_{test}$)가 성능에 미치는 영향을 분석했다. 가치 함수가 잘 일반화된 경우 추론 시 단계를 늘릴수록 성능이 향상되는 '테스트 타임 스케일링' 현상이 관찰되었으며, 가치 함수의 외삽 오류가 큰 경우 단계를 줄여 참조 분포에 가깝게 유지함으로써 안정성을 확보할 수 있음을 보여주었다.

VGF는 정책 파라미터 $\theta$에 대한 최적화 대신 입자 집합 $\{a_i\}_{i=1}^N$의 역학을 직접 다룬다. 이는 정책 그래디언트 방법론에서 발생하는 높은 분산 문제를 회피하며, 특히 다봉 분포(Multimodal Distribution)를 캡처하는 데 유리하다. 이론적으로 VGF를 통해 생성된 입자 분포와 참조 분포 사이의 MMD(Maximum Mean Discrepancy) 거리가 운송 예산에 의해 유계(Bounded)됨을 증명하여 암묵적 규제의 효과를 수학적으로 뒷받침했다.

구현 측면에서 VGF는 가치 함수의 액션에 대한 그래디언트 $abla_a Q(s, a)$를 효율적으로 계산하기 위해 별도의 그래디언트 네트워크 $f(s, a)$를 학습시켜 추론 속도를 가속화할 수 있다. 또한 MaxEnt(Maximum Entropy) 항을 선택적으로 제거함으로써 결정론적인 최적화와 확률적인 탐색 사이의 균형을 조절할 수 있는 구조를 갖추고 있다.

참조 분포가 극도로 편향되어 있거나 최적 행동 영역과 전혀 겹치지 않는 경우 성능이 제한될 수 있다. 또한 가치 함수의 그래디언트에 의존하기 때문에 가치 함수 자체가 심하게 잘못 학습된 경우 입자가 엉뚱한 방향으로 이동할 위험이 있으며, 이를 해결하기 위한 분포 재가중(Distribution Reweighting) 기법과의 통합이 향후 과제로 남아 있다.