HMC의 운동량 다시 뽑기: 열원과 맞닿는 순간
HMC는 목표 분포 π(q)를 e^(−U(q))로 적고, 갱신할 때마다 운동량을 표준정규분포에서 새로 뽑는다. 위치와 운동량을 함께 가진 계의 볼츠만 분포가 곱으로 쪼개진다는 것을 알았으니, 이제 이 두 단계를 물리의 말로 읽을 수 있을까? 특히 매번 운동량을 새로 뽑는 단계는 계에 무엇을 해 주는 것일까?
ML에서: 등고선 위와 등고선 사이
HMC가 목표 분포를 위치에너지 U(q) = −ln π(q)로 옮기는 것은 π를 온도 1의 볼츠만 분포로 읽는 일이다. 사후분포처럼 정규화 상수를 모르는 π라도 괜찮은 이유가 여기서 보인다. 정규화 상수는 U에 더해지는 상수, 곧 에너지의 기준점일 뿐이고, HMC는 수락 확률 e^(−Δℋ)에서 에너지의 차이만 쓴다. 운동량을 표준정규분포에서 뽑는 것은 질량 1인 입자의 운동량을 온도 1의 맥스웰 분포에서 뽑는 것이고, 결합 분포(joint) e^(−ℋ)가 위치와 운동량의 곱으로 쪼개지므로 운동량을 버리면 π가 남는다.
이제 한 번의 갱신을 열원의 언어로 읽을 수 있다. HMC가 궤적을 따라가는 leapfrog(운동량 반 걸음, 위치 한 걸음, 운동량 반 걸음을 되풀이하는 적분법) 궤적은 에너지를 거의 보존하므로 한 등고선 위를 돌고, 운동량을 새로 뽑는 순간 계는 온도 1인 열원과 에너지를 주고받아 다른 등고선으로 옮겨 간다. 베탄코트(Betancourt, 2017)의 HMC 해설은 결합 분포를 등고선 위의 미시정준 분포와 에너지 자체의 분포로 나눈 뒤 이렇게 적었다. 「해밀턴 마르코프 사슬 전체가 두 단계로 나뉜다. 개별 등고선의 결정론적 탐색과, 등고선 사이의 확률적 탐색이다.」 이 장을 마친 독자는 이것을 「궤적은 전체 에너지가 고정된 계의 균등 분포를 돌아다니고, 운동량 다시 뽑기는 열원과 맞닿아 에너지를 볼츠만 분포대로 바꾼다」로 읽는다.
문제 12. 운동량을 다시 뽑지 않는 HMC
목표 분포가 1차원 표준정규분포(U(q) = q²/2)이다. 민준은 계산을 아끼려고 처음에 한 번만 운동량을 뽑고(q = 1, p = 0.5), 그 뒤로는 운동량을 다시 뽑지 않은 채 leapfrog 궤적(걸음 크기 0.15)을 계속 따라가며 10걸음마다 q를 기록했다. (가) 기록된 q는 표준정규분포를 따르는가? (나) 서연은 더 멀리 가려고 운동량을 N(0, 4)에서 뽑되, ℋ는 그대로 q²/2 + p²/2로 계산해 수락 여부를 정했다. q의 분산은 얼마가 되는가?
따를 것 같아요. 해밀턴 흐름은 e^(−ℋ)를 바꾸지 않으니까요. 목표 분포를 흔들지 않는 흐름을 따라가며 기록했으니 그대로 표본이죠. 수락 확률도 거의 1이고, 운동량 뽑는 계산도 아꼈고요.
돌려 봤는데 이상해. 2만 개를 기록했는데 |q|가 1.12를 한 번도 안 넘어. 표준정규분포면 1.12를 넘는 게 26%는 돼야 하는데. 분산도 1이 아니라 0.626이야.
어? 버그인가… 코드는 맞는데.
처음 상태에서 ℋ를 계산해 봐요.
0.5 + 0.125 = 0.625요. 그리고 ℋ는 보존되니까… 아, 상태는 ℋ = 0.625인 원 위만 돌아요. 반지름이 √1.25 = 1.118이라 그 밖으로는 절대 못 나가요.
분산이 0.626인 것도 설명돼. 원 위에서는 q²/2와 p²/2의 시간 평균이 같아서 각각 ℋ/2니까 ⟨q²⟩ = ℋ = 0.625야. 게다가 원을 도는 동안 q가 오래 머무는 곳은 속도가 느린 양 끝 ±1.118 근처야. 봉우리가 가운데가 아니라 양 끝에 있는 분포라서 정규분포랑 정반대네.
민준 학생이 말한 「흐름이 e^(−ℋ)를 바꾸지 않는다」는 맞아요. 그런데 그건 처음부터 e^(−ℋ)로 퍼져 있는 상태들 전체의 이야기고, 궤적 하나는 등고선 하나에 갇혀 있어요. 민준 학생이 뽑은 건 그 등고선 위의 미시정준 분포예요. e^(−ℋ)는 여러 등고선을 볼츠만 인자의 비중대로 섞은 것이라서, 등고선 사이를 옮겨 다니는 단계가 따로 있어야 해요.
그게 운동량을 다시 뽑는 단계군요. 새 운동량을 뽑으면 ℋ가 바뀌니까요. 계에 에너지를 넣었다 뺐다 하는 열원 노릇이에요.
그래요. 그럼 그 열원의 온도는 얼마죠?
운동량을 N(0, 1)에서 뽑는 건 e^(−p²/2)에서 뽑는 거니까, 운동에너지 p²/2의 볼츠만 분포를 온도 1로 뽑은 거예요. 목표도 e^(−U)니까 온도 1이고, 맞아떨어지네요. 그럼 (나)처럼 N(0, 4)에서 뽑으면 더 멀리 가니까 더 좋은 거 아니에요?
나도 그렇게 생각하고 돌렸는데… 걸음 크기 0.15, 10걸음으로 20만 번 갱신했더니 q의 분산이 3.95야. 수락률은 99%인데. 1이 나와야 하는데 거의 4가 나왔어.
문제 10을 떠올려 봐요. 분산이 4인 정규분포는 어떤 온도의 볼츠만 분포였죠?
e^(−p²/(2·4))니까 온도 4요. 아… 운동량은 온도 4인 열원에서 뽑고, 수락 규칙은 온도 1로 계산했어요. 수락률이 거의 1이라 수락 단계는 사실상 아무 일도 안 했고, 위치는 온도 4인 열원에 맞닿아 있었던 거예요. 걸음 크기를 0.05로 줄이고 걸음 수를 서른으로 늘렸더니 분산이 3.99까지 가요. 위치의 분산이 운동량을 뽑은 온도를 그대로 따라가요.
두 사람이 한 가지씩 짚었어요. 그리고 운동량 분포와 운동에너지는 따로 정할 수 있는 게 아니에요. U = −ln π로 위치에너지를 정했듯이, 운동에너지도 K = −ln(운동량 분포)로 정해야 해요. N(0, 4)에서 뽑고 싶으면 K = p²/8, 곧 질량을 4로 두면 돼요.
K = p²/8로 바꿔서 같은 설정으로 돌렸더니 분산 0.99요. 제자리로 왔어요. 열원이 두 군데서 들어오면 온도를 하나로 맞춰야 하는 거네요.
그러고 보니 HMC가 운동량을 표준정규분포에서 뽑는 건 편해서 고른 게 아니었네요. 운동에너지를 p²/2로 정한 순간 이미 정해진 거였어요.