Optimization + MCMC methods: Langevin Dynamics, Langevin MCMC

 

아래 references의 강의 및 자료들을 공부하고 짧게 정리한 내용입니다! 저도 공부하면서 정리한 내용이라 틀린 것이 있다면 언제든 댓글 달아주신다면 감사하겠습니다~ :D

____

 

1. Optimization + MCMC 결합 이유

 

그림과 같이 target distribution P(x), MCMC에서 sample proposal을 생성하는 간단한 gaussian Q(x'|x)가 있다. (MCMC 설명: 이전 포스팅 참조) 이 때 Q(x'|x)가 gaussian이므로 어떤 방향으로든지 이동할 확률이 같아 결과적으로 랜덤하게 이동하게 되는데, 이를 Random Walk이라고 한다. 물론 잘못된 sample이 뽑힌 경우 accept rate를 조절하여 해결할 수 있지만, 수렴 속도가 매우 느려지고 오랜 기간동안 낮은 accept rate를 유지해야 할 수도 있다. 특히 Q의 variance가 클수록 exploration이 잘 되는 대신 이러한 문제가 심해진다. Q의 variance를 낮춰서 이 문제를 해결하려고 하면 sample간의 correlation이 심해지고, 결과적으로 효율적으로 sample을 뽑지 못하게 된다. (Q의 variance에 관한 설명: 이전 포스팅 참조)

 

따라서 만약 p(x) distribution의 gradient 정보를 이용할 수 있다면, Random Walk 대신 훨씬 더 효율적으로 p(x)에서 likelihood가 높은 방향으로 이동할 수 있다. 이러한 이유로 MCMC에 optimization concept을 적용하게 되었고, 특히 score-based model을 이용하게 된다.

 

2. Langevin Dynamics

Stochastic Differential Equation (SDE)는 시간에 따른 어떤 variable의 변화량을 나타내기 위한 미분방정식이다. 특이한 점은 경향성을 갖고 deterministic하게 움직이는 term과 random하게 움직이는 term 두가지로 구성되어 있다는 것이다. 따라서 완전히 deterministic한 것 뿐만 아니라 stochastic한 것도 모델링하므로 실제 세계를 모델링하는 데 더 적합하다.

 

 

위 식은 SDE를 나타내는 식이다. SDE를 풀기 위해 Euler-Maruyama method를 이용해서 시간을 discretize하면 아래와 같이 나타낼 수 있다: (SDE, Euler-Maruyama에 대한 설명은 이 블로그에 잘 되어있습니다)

 

 

여기서 갑자기 U라는 함수가 나오는데, 이는 물리학에서 분자 운동을 나타내기 위한 에너지 표현을 이용한 것이다. 대략적으로 Hamiltonian Dynamics에서는 분자가 가진 총 에너지를 H(p, x) = K(p) + U(x)로 나타낸다고 한다. 이 때 K(p)는 kinetic energy로 p는 momentum을, U(x)는 potential energy로 x는 position을 의미한다.

 

간단하게 핵 근처를 돌고있는 전자를 생각하면, 전자의 에너지는 위치에너지 + 운동에너지로 나타낼 수 있다. 여기서도 이와 유사하게 K(p) + U(x)의 형태로 식을 표현한 것이다. 이 때 목적은 x의 위치를 구하기 위한 것이므로 U(x)가 deterministic한 term이고, K(p)는 stochastic한 term으로 p는 N(0, I)에서 sampling하여 구한다고 한다.

 

결과적으로 위 식의 두번째 줄은 이러한 분자의 에너지 함수를 대입한 것으로, U는 position, N(0, I)는 momentum을 나타낸다.

이 식을 Langevin Equation이라고 하며, Langevin Dynamics를 나타낸 식이다.

사실상 SDE를 풀기 위해 Euler-Maruyama로 discretize한 식에 분자 운동을 나타내는 energy 함수만 대입한 것으로 볼 수 있다.

갑자기 분자운동이 나오는데, 결과적으로 U는 deterministic한 term, N(0,I)는 randomness를 나타내는 term이라고 생각하면 된다. MCMC에서 gradient를 이용해 큰 방향성을 유지하면서 (deterministic) greedy하지 않도록 exploration하는 (stochastic) 행동 방식을 모델링하기에 매우 유사하기 때문에 가져다 쓴 것이라고 생각하면 된다.

Score Function에 관한 식으로 변경.

최종적으로 Langevin Equation에 ∇U(x) 대신 -∇log p(x)를 대입할 것이다. 그 이유는 다음과 같다:

 

첫 번째 식은 distribution p(x)를 energy-based representation으로 나타낸 것이다. 이 식에 양변에 log를 취하고, 미분을 하면 두 번째와 같은 식이 된다. 결국 ∇U(x) = -∇log p(x)이다. 또한 Langevin Equation에 sigma 대신 root 2를 대입하는데, 이렇게 수식을 변경하는 이유는 크게 두 가지가 있다:

• Score Function ∇log p(x)에 관한 식으로 변경하기 위해서
• 이렇게 대입하면 결국 Langevin Dynamics의 Stationary Distribution이 Target Distribution이 되기 때문에

 

두 번째 이유의 경우, 이 dynamics를 따라 iterative하게 update했을 때 수렴하는 distribution (stationary distribution)이 결국 우리가 구하고자 하는 target distribution이 된다는 뜻이다. (이에 대한 증명은 아래 reference의 맨 마지막 강의나, 이 블로그를 참조하면 된다) 결국 아래 식을 계속 update하면 stationary distribution으로 수렴하고, 그게 우리가 원하는 target distribution이 된다는 뜻이다:

 

3. Langevin MCMC

결과적으로 위에서 구한 update 식을 이용하여 수렴할 때까지 iterative하게 sample을 업데이트 하면 된다:

 

위의 알고리즘은 Unadjusted Langevin MCMC로, 항상 sample을 accept하는 경우이다. 물론 MCMC처럼 accept rate를 설정하는 경우도 있다.

 

4. MCMC가 굳이 필요한 이유?

그렇다면 그냥 optimization 하면 되지 왜 굳이 MCMC를 하는 것일까?

 

먼저 현재 추정하는 것은 classification과 같이 conditional distribution p(y|x)가 아닌 어떤 distribution p(x)이기 때문에, 더 나타내기가 어렵다. 여러가지 방법이 있지만 예를 들어 hidden variable z를 이용하여 p(x)를 표현하는 경우, hidden variable z에 대해 marginalize 해야하기 때문에 식을 연산할 수 없고 다른 방법이 필요하다. 이 때 variational inference나 MCMC 등의 방법을 쓸 수 있다. 또한 만약 p(x) = a(x)b(x)c(x) 처럼 p(x)가 여러 distribution의 AND operation이고, 이 식이 closed form으로 연산이 불가능하다면 MCMC를 쓸 수 있다. (관련 논문 참조)

 

두 번째로, optimized된 model로부터 high probability를 가지면서 다양한 sample을 생성하고 싶을 때 활용할 수 있다. 예를 들면 LLM이나 diffusion model을 inference할 때, 더 다양하고 품질 높은 답변을 생성하기 위해 이런 방법을 사용할 수 있을 것이다.

 

 

[References]

https://www.cs.cmu.edu/~epxing/Class/10708-20/lectures.html

10-708 – Lectures (tentative)

https://www.youtube.com/playlist?list=PLoROMvodv4rPOWA-omMM6STXaWW4FvJT8

Stanford CS236: Deep Generative Models I 2023 I Stefano Ermon

https://www.youtube.com/watch?v=3-KzIjoFJy4