자기회귀로 다시 쓴 볼츠만 생성기, 플로우의 한계를 넘어선 분자 샘플링의 전환점

분자가 열역학적 평형에서 어떤 형태로 존재하는지를 빠르게 알아내는 일은 통계물리학의 오래된 난제다. 단백질 하나가 접히는 과정, 약물 분자가 표적에 달라붙는 순간의 자세, 이런 것들은 모두 볼츠만 분포라는 확률법칙을 따른다. 에너지가 낮은 상태일수록 더 자주 관측되지만, 그 사이를 가로막는 높은 에너지 장벽 때문에 전통적인 분자동역학 시뮬레이션은 한 상태에서 다른 상태로 넘어가는 데 막대한 계산 시간을 쏟아붓는다. 게다가 연속한 시뮬레이션 스냅숏들은 서로 강하게 상관되어 있어, 통계적으로 독립적인 표본을 얻기란 더욱 어렵다.

볼츠만 생성기(Boltzmann Generator)는 이 문제에 대한 우아한 우회로였다. 생성모델로 분자 배치를 단숨에 만들어내되, 그 모델이 각 표본의 정확한 확률(likelihood)을 함께 내놓을 수 있다면, 중요도 샘플링(importance sampling)이라는 보정 절차를 통해 편향 없는 평형 표본을 복원할 수 있다. 핵심은 '정확한 확률을 계산할 수 있는 생성모델'이라는 조건이었고, 지금까지 이 자리를 거의 독점해온 것이 정규화 플로우(normalizing flow)였다.

플로우라는 황금 새장

정규화 플로우는 단순한 기준 분포를 가역적인 변환으로 비틀어 복잡한 목표 분포로 옮긴다. 변환이 가역적이어야 한다는 제약 덕분에 확률을 정확히 계산할 수 있지만, 바로 그 제약이 족쇄가 된다. 이산시간 플로우, 즉 결합층(coupling layer)을 쌓는 방식은 가역성을 보장하기 위해 표현력을 희생한다. 반대로 연속시간 플로우는 미분방정식을 적분해야 확률을 얻기 때문에 표현력은 풍부하지만 계산 비용이 가파르게 치솟는다. 더 근본적으로, 가역 변환은 위상(topology)을 보존하기 때문에 매끈한 기준 분포에서 여러 골짜기로 갈라진 다봉형 목표 분포로 부드럽게 이어 붙이기가 구조적으로 까다롭다. 분자가 커질수록 이 한계는 더 도드라진다.

이 논문이 제안하는 자기회귀 볼츠만 생성기(ArBG)는 플로우라는 패러다임 자체에서 걸어 나온다. 발상은 의외로 익숙하다. 언어모델이 문장을 한 단어씩 이어 쓰듯, 분자의 좌표를 하나씩 순차적으로 예측하는 것이다. 전체 확률은 조건부 확률들의 곱이라는 확률의 연쇄법칙으로 자연스럽게 분해되므로, 야코비 행렬식 계산도 가역성 제약도 필요 없이 정확한 확률이 손에 들어온다. 가역성이라는 위상학적 굴레가 사라지자 표현력의 천장도 함께 사라졌다.

언어모델의 구조를 빌려오다

자기회귀 방식의 또 다른 이점은 생성을 한 단계씩 진행하는 동안 중간에 개입할 수 있다는 점이다. 추론 시점에 순차적으로 조건을 끼워 넣어 원하는 방향으로 표본을 유도할 수 있다. 무엇보다 이 틀은 트랜스포머처럼 대규모 언어모델에서 검증된 구조를 그대로 가져다 쓸 수 있어, 모델과 데이터를 키웠을 때의 확장성이 플로우보다 훨씬 유리하다.

실험 결과는 이 방향이 빈말이 아님을 보여준다. ArBG는 모든 벤치마크에서 플로우 기반 모델을 앞섰고, 특히 10잔기 펩타이드 시그놀린(Chignolin)처럼 덩치가 큰 계에서 격차가 두드러졌다. 연구진은 한 걸음 더 나아가 1억 3200만 파라미터 규모의 전이 가능 모델 Robin을 내놓았다. Robin은 학습에 쓰이지 않은 8잔기 시스템에 대해서도 별도 미세조정 없이 표본을 생성하는 제로샷 설정에서 종전 최고 성능을 갈아치웠고, 에너지 오차 지표 E-W2를 60% 넘게 줄였다.

이 연구가 시사하는 바는 단순히 벤치마크 숫자 하나가 좋아졌다는 데 그치지 않는다. 분자 과학을 위한 생성모델이 그동안 플로우의 수학적 우아함에 묶여 있었다면, ArBG는 언어모델 시대에 축적된 확장성의 노하우를 분자 샘플링으로 끌어들이는 통로를 연다. 단백질과 펩타이드처럼 자유도가 큰 계로 갈수록 이 통로의 가치는 커질 것이고, 전이 가능한 기반모델이라는 Robin의 방향성은 분자 시뮬레이션이 거대 모델의 문법을 받아들이기 시작했다는 신호로 읽힌다.