리우빌 방정식: 흐름을 따라가면 밀도는 그대로
넓이가 변하지 않는다는 것은 상태들의 무리를 부피로 본 이야기다. 그렇다면 상태들이 위상공간에 확률밀도 ρ(q, p)로 퍼져 있을 때, 그 밀도는 흐름을 따라 어떻게 변할까?
따라가는 관찰자와 서 있는 관찰자
진자의 넓이 0.04짜리 정사각형(각도 2.4~2.6, 운동량 0~0.2)에 상태 1000개를 고르게 뿌려 보자. 밀도는 넓이 하나당 들어 있는 상태의 수이므로 처음에는 1000 ÷ 0.04 = 25,000이다. 이 1000개를 흐름을 따라가며 지켜보면, 상태들은 무리 밖으로 빠져나가지 않으니 개수는 1000 그대로이고, 리우빌 정리 때문에 무리가 차지하는 넓이도 0.04 그대로다. 그래서 시간 5 뒤에도, 10 뒤에도 무리를 따라가는 관찰자가 재는 밀도는 25,000이다. 그런데 처음 정사각형 자리에 창을 고정해 두고 그 안의 상태를 세면 이야기가 다르다. 시간 5 뒤 무리는 위상공간의 반대쪽(각도 −2.5 근처)으로 흘러가 있어서 창 안에는 상태가 하나도 없다(python으로 흘려 본 값).
같은 흐름을 두고 따라가는 관찰자는 「밀도가 그대로」라 하고, 서 있는 관찰자는 「밀도가 바뀐다」고 한다. 둘은 모순이 아니다. 서 있는 자리의 밀도가 바뀌는 것은 다른 밀도를 가진 상태들이 흘러 들어오고 나가기 때문이다. 「따라가면 그대로」를 한 자리에 서서 보는 밀도의 변화로 옮겨 적으면 다음과 같고, 이를 리우빌 방정식 (Liouville equation)이라 한다.
왼쪽 첫 항 ∂ρ/∂t는 서 있는 관찰자가 보는 밀도의 변화율이고, 뒤의 두 항은 상태들이 흘러 지나가며 실어 오는 밀도의 차이다. 셋을 합한 것이 흐름을 따라 움직이며 잰 밀도의 변화율이므로, 이 식은 「흐름을 따라가면 밀도가 변하지 않는다」는 말을 그대로 옮긴 것이다. 여기서 쓸모 있는 결론이 하나 나온다. 밀도가 에너지만의 함수, 이를테면 ρ ∝ e^(−ℋ)이면 서 있는 관찰자가 보아도 밀도는 시간이 지나도 변하지 않는다. 각 상태는 자기 등고선을 벗어나지 않고, 같은 등고선 위의 밀도는 어디서나 같아서, 흘러 들어오는 상태와 흘러 나가는 상태의 밀도가 같기 때문이다. 해밀턴 흐름은 이런 분포를 흔들지 않고 그 안에서 상태들만 옮겨 놓는다.
ML에서: 연속 normalizing flow의 로그밀도 식
normalizing flow는 표준정규분포처럼 뽑기 쉬운 분포의 샘플을 되돌릴 수 있는 변환으로 옮겨 복잡한 데이터 분포를 만드는 생성 모델이다. 그 가운데 연속 normalizing flow는 표준정규분포에서 뽑은 샘플을 신경망이 정한 속도장 dx/dt = f(x, t)를 따라 흘려보내 데이터 분포로 옮긴다. 학습하려면 샘플의 로그밀도가 흐름을 따라 어떻게 변하는지 알아야 하는데, Chen 외(2018)의 Neural ODE 논문은 이를 「순간 변수 변환」 정리로 적었다.
이 식은 리우빌 방정식을 이끌어 낸 논리와 같다. 흐름을 따라가는 작은 영역의 확률은 그대로이고 넓이는 발산의 빠르기로 늘어나므로, 밀도의 로그는 발산만큼 줄어든다. 해밀턴 흐름에서는 발산이 0이라 밀도가 변하지 않았던 것이고, 신경망 속도장은 발산이 0이 아니어서 분포의 모양을 바꿀 수 있는 것이다. 식의 야코비 행렬(Jacobian matrix)은 속도의 각 성분을 각 좌표로 미분해 늘어놓은 표이고, 그 대각선을 더한 대각합(trace)이 곧 발산이다. 층을 쌓는 이산 flow는 층마다 야코비 행렬식(행렬식은 그 변환이 작은 영역의 넓이를 몇 배로 바꾸는지를 나타낸다)의 로그를 더해야 하는데, 연속 flow에서는 이 계산이 무엇으로 바뀌는지 아래 문제 9에서 셈해 본다. Grathwohl 외(2019)의 FFJORD는 이 방식을 고차원 데이터까지 늘렸다.
문제 8. 에너지만의 함수가 아닌 밀도
그네(ℋ = p²/2 + q²/2)의 위상공간에 밀도 ρ(q, p) ∝ e^(−(q² + 4p²)/2)로 상태들을 뿌렸다(위치 방향으로 넓고 운동량 방향으로 좁은 타원 모양). 점 (q, p) = (1, 1)에 서서 보면 밀도는 늘어나는가, 줄어드는가? 시간이 π/2 지난 뒤 분포는 어떤 모양인가?

리우빌 방정식이 밀도는 안 변한다는 식이잖아요. 그러니까 그대로요.

그 식에서 「안 변한다」는 누가 잰 밀도였죠?

흐름을 따라가는 쪽이요… 이 문제는 한 점에 서서 보는 거네요.

그럼 리우빌 방정식을 ∂ρ/∂t = −(q̇ ∂ρ/∂q + ṗ ∂ρ/∂p)로 옮겨서 넣으면 돼. q̇ = p, ṗ = −q이고, ∂ρ/∂q = −qρ, ∂ρ/∂p = −4pρ니까 ∂ρ/∂t = −(−qp + 4qp)ρ = −3qpρ. (1, 1)에서는 −3ρ라서 줄어들어.

(1, −1)에 서 있으면 +3ρ라서 거꾸로 늘어나네요. 그네 흐름은 원을 따라 시계 방향으로 도니까, 타원 전체가 빙글 도는 거예요.

그래요. π/2만큼 돌면 위치 방향으로 좁고 운동량 방향으로 넓은 타원, 곧 e^(−(4q² + p²)/2) 모양이 돼요. 그럼 밀도가 어떤 모양이면 어디에 서 있어도 변하지 않을까요?

−3qp가 나온 건 두 방향의 계수 1과 4가 달라서예요. 둘이 같으면, 곧 q² + p², 에너지만의 함수면 두 항이 지워져서 0이 돼요. 원 모양 분포는 돌려도 그대로니까요.

동그란 피자는 돌려도 모양이 그대로인데 길쭉한 빵은 돌리면 방향이 바뀌는 거랑 같네요.
문제 9. 연속 흐름의 로그밀도
속도장이 dx/dt = Ax, A = [[−1, 2], [−2, −1]]인 2차원 연속 흐름을 따라 샘플 하나를 시간 1만큼 흘려보냈다. 그 샘플에서의 로그밀도는 얼마나 변하는가?

normalizing flow는 로그밀도에서 log|det J|를 빼잖아요. det A = 1 + 4 = 5니까 ln 5 ≈ 1.609만큼 줄어요.

그건 층 하나짜리 변환 x ↦ Ax일 때 공식이잖아. 여기서 A는 점을 옮기는 변환이 아니라 속도장이야. 순간 변수 변환 식은 −tr(∂f/∂x) = −tr A = 2니까, 시간 1 동안 로그밀도가 2만큼 늘어나.

줄어드는 게 아니라 늘어난다고? 부호가 반대인데 어떻게 확인하지?

이산 공식으로 검산할 수 있어요. 시간 1 뒤의 변환은 행렬 지수 e^A이고, 그 행렬식을 구하면 돼요.

선형대수에서 배운 공식 det e^A = e^(tr A)를 쓰면 e^(−2) ≈ 0.135예요. 넓이가 0.135배로 줄어드니까 밀도는 그만큼 진해지고, 로그밀도는 −ln e^(−2) = 2만큼 늘어요. 두 방법이 같은 답을 주네요. 대각합 공식이 행렬식 공식의 연속 버전이었어요.

A의 대각선이 둘 다 −1이라 수축하는 흐름이었네요. 그럼 대각선 밖의 2와 −2는요?

그 부분만 떼어 보면 (2x₂, −2x₁)인데, ℋ = x₁² + x₂²로 두고 x₁을 위치, x₂를 운동량으로 보면 정확히 해밀턴 방정식이에요. 샘플을 원을 따라 돌리기만 하고 대각합은 0이라 밀도를 바꾸지 못해요. 밀도를 바꾸는 건 발산이 0이 아닌 부분뿐이에요.

그러니까 행렬식 5는 회전과 수축을 한꺼번에 센 숫자라서 여기서 쓸 게 아니었고, 대각합은 회전을 저절로 걸러 주는 거네요.

그럼 속도장이 신경망이라 행렬 지수를 손으로 구할 수 없을 때, 구현에서는 매 순간 무엇을 계산하면 될까요?

속도장을 좌표로 미분한 표의 대각선 성분만 더해서 시간으로 쌓으면 되겠네요. 층을 쌓는 flow처럼 층마다 큰 행렬의 행렬식을 구하지 않아도 되고요.

그래요. 대각선만 필요하니 계산이 훨씬 싸요. 차원이 아주 높으면 그 대각합마저 무작위 벡터로 어림하기도 하고요.