One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective

기술 요약: One More Time: 강화 학습 관점에서 재조명한 신경 양자 상태(Neural Quantum States)

문제 정의

신경 양자 상태(NQS)는 특히 Born 분포로부터 정확하고 독립적인 샘플링을 가능하게 하는 자기회귀(autoregressive) 모델을 활용할 때, 양자 다체 파동함수를 근사하기 위한 확장 가능한 프레임워크를 제공한다. 그러나 이러한 모델의 최적화는 여전히 중요한 병목 현상으로 남아 있다. 기존 방법들은 다음과 같은 극명한 트레이드오프에 직면해 있다:

• 1차 최적화 도구 (예: Adam): 계산 효율적이고 확장성이 뛰어나지만, 변분 파동함수의 기하학적 구조를 무시하여 NQS 응용 분야에서 불안정하거나 부정확한 수렴을 초래하는 경우가 많다.
• Stochastic Reconfiguration (SR) 및 변형 알고리즘 (예: minSR): 기하학적으로 원칙적이지만(자연 경사 하강법 근사), 거대 네트워크의 경우 크고 종종 조건이 나쁜(ill-conditioned) 선형 시스템을 풀어야 하므로 계산 비용이 매우 높다. 또한 복잡한 체제(regime)에서 수치적 취약성을 겪는다.

자기회귀 NQS는 샘플링의 병목 현상을 해결하지만, 대규모 자기회귀 파동함수 모델을 위한 안정적이고 확장 가능한 최적화는 여전히 미해결 과제로 남아 있다.

방법론

저자들은 강화 학습(RL)의 근접 정책 최적화(Proximal Policy Optimization, PPO)에서 영감을 받은 신뢰 영역(trust-region) 알고리즘인 **근접 파동함수 최적화(Proximal Wavefunction Optimization, PWO)**를 제안한다. 핵심 방법론은 변분 에너지 최소화와 RL의 정책 경사(policy-gradient) 최적화 사이의 공식적인 등가성에 기초한다.

1. RL-VMC 등가성: 본 논문은 stoquastic 해밀토니안에 대해, 변분 에너지의 기울기가 Born 분포에 대한 어드밴티지 가중 정책 경사(advantage-weighted policy-gradient) 업데이트로 표현될 수 있음을 입증한다. 이 매핑에서:

   • 스핀 구성은 행동(action)에 대응한다.
   • 중심화된 국소 에너지($E_{loc}(s) - \mathbb{E}[E_{loc}]$)는 어드밴티지(advantage)에 대응한다.
   • Born 분포는 정책(policy)에 대응한다.

2. 이중 채널 최적화: 단일 확률 분포를 다루는 표준 PPO와 달리, NQS 파동함수는 복소수 값(complex-valued)을 가지므로 진폭(amplitude)과 위상(phase)을 별도로 처리해야 한다. PWO는 두 가지 클리핑된 대리 목적 함수(clipped surrogate objectives)를 도입한다:

   • 진폭 채널 (Amplitude Channel): 확률 비율($r_\theta = P_\theta / P_{\theta_{old}}$)을 사용하여 업데이트를 제한하며, 이는 표준 PPO를 모방한다. 이를 통해 업데이트된 파동함수가 확률 질량 측면에서 참조 모델과 가깝게 유지되도록 보장한다.
   • 위상 채널 (Phase Channel): 확률 비율 대신 *래핑된 위상 증분(wrapped phase increment, $\phi_\theta$)*을 제한하는 새로운 클리핑된 대리 함수를 도입한다. 이는 변분 몬테카를로(VMC) 기울기의 허수 부분을 따르면서도 파동함수의 급격한 회전을 방지한다.

3. 알고리즘 구현: PWO는 참조 분포 $P_{\theta_{old}}$에서 추출된 고정된 배치(batch) 상에서 여러 번의 내부 업데이트($K$)를 수행한다. 이들은 분포 불일치를 보정하기 위해 중요도 가중치(importance weighting)와 stop-gradient 연산자를 사용하여 샘플을 재사용한다. 이는 명시적인 행렬 역행렬 계산을 피함으로써 1차 최적화 도구의 확장성을 유지하면서도 신뢰 영역 제약을 부과한다.

주요 기여

• 공식적 연결: 변분 에너지 최소화를 어드밴티지 정책 경사 RL로 공식적으로 연결하여, NQS의 기울기가 Born 분포에 대한 어드밴티지 가중 형태를 가짐을 입증하였다.
• PWO 알고리즘: 확률 비율의 변화를 클리핑하고 위상 증분을 클리핑하는 진폭 채널과 위상 채널을 갖춘 근접 파동함수 최적화(PWO)를 도입하였다.
• 이론적 보장: PWO 대리 함수가 신뢰 영역 개선 경계(improvement bound)를 만족하며, 온폴리시(on-policy) 상태일 때 1차 업데이트가 정확한 VMC 기울기 방향과 일치함을 증명하였다.
• 확장성 입증: 15억 개의 파라미터를 가진 RWKV-7 대규모 언어 모델(LLM)을 자기회귀 NQS로서 성공적으로 미세 조정함으로써, 이전 연구보다 3개 차수(orders of magnitude) 이상 큰 규모에서의 최적화를 입증하였다.

실험 결과

저자들은 Ising 및 J1–J2 모델을 포함한 여러 벤치마크에서 PWO를 Adam, minSR, SPRING과 비교 평가하였다:

• 1D 스핀 시스템 (Ising 및 J1–J2):
  • Ising 모델: PWO는 약 5분 만에 상대 오차 $10^{-7}$에 도달하며, 이는 minSR(~30분) 및 Adam보다 훨씬 빠르다.
  • 좌절된(Frustrated) J1–J2 체인: PWO는 감소하는 V-score와 함께 안정적인 수렴을 유지했다. 반면, minSR은 수치적으로 불안정했으며(10번의 실행 중 6번에서 NaN 발생), Adam은 훨씬 높은 오차 수준에서 정체되었다.
• 2D 좌절된 J1–J2 모델: $10 \times 10$ 격자에서, PWO는 Adam보다 더 빠르게 낮은 에너지 영역에 도달했으며, 동일한 실제 시간(wall-clock) 예산 내에서 더 낮은 분산 기반 오차 신호(V-score)를 유지했다.
• 확장성: 모델 크기가 증가함에 따라 ("tiny"에서 "medium" NQS까지), PWO는 일관된 안정성과 단조로운 에너지 개선을 보여주었다. Adam은 더 큰 모델에서 안정성과 효율성이 저하되었고, minSR은 경쟁력 있는 에너지 값에 도달하지 못했다.
• 대규모 미세 조정: PWO로 미세 조정된 1.5B 파라미터 RWKV-7 모델은 Adam보다 더 낮은 최종 상대 오차와 V-score를 달면서 전체 예산 동안 안정적인 훈련을 유지했다.

의의 및 주장

본 논문은 PWO가 RL의 신뢰 영역 원리를 활용함으로써 1차 최적화 도구의 확장성과 기하학적 방법(SR 등)의 안정성 사이의 간극을 메운다고 주장한다.

• 어려운 체제에서의 안정성: 저자들은 PPO 스타일의 근접 목적 함수가 전통적인 방법들이 실패하거나 불안정해지는 어려운 좌절된(frustrated) 양자 해밀토니안에서 표현력이 풍부한 NQS를 훈련하는 데 필요한 안정성을 제공한다고 단언한다.
• 샘플 효율성: 명시적인 행렬 역행렬 계산 없이 여러 업데이트에 걸쳐 샘플을 재사용할 수 있게 함으로써, PWO는 실제 시간 기준의 수렴 속도를 향상시킨다.
• 규모(Scale): 1.5B 파라미터 모델에 대한 성공적인 적용은 NQS 최적화가 현대의 대규모 시퀀스 모델링 영역으로 진입할 수 있음을 시사한다. 저자들은 이것이 규모 자체가 현재의 장벽을 극복할 수 있는지 조사하는 문을 열어줄 것이라고 상정하며, 이는 AI의 "고통스러운 교훈(bitter lesson)"과 궤를 같이한다.

이 연구는 LLM이 1D 해밀토니안을 위한 최적의 귀납적 편향(inductive bias)이라고 주장하는 것이 아니라, 전례 없는 규모에서 근접 목적 함수의 효용성을 입증하기 위해 RWKV-7 모델을 사용한 것이다.