🤖 machine learning

Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps

이 논문은 F2D2와 같은 기존 베이스라인에 비해 추론 속도를 최대 10배까지 높이면서도 훈련 분산과 시간을 크게 줄인, 흐름 기반 생성 모델을 위한 확장 가능하고 Hutchinson-free한 가능도 증류 방법인 SCALLOP을 소개한다.

원저자: RuiKang OuYang, Hanlin Yu, Xinyue Ai, Yutong He, Nicholas M. Boffi, Pradeep Ravikumar, Jose Miguel Hernandez-Lobato, Max Simchowitz, Benjamin Kurt Miller, Omar Chehab

게시일 2026-06-30

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: RuiKang OuYang, Hanlin Yu, Xinyue Ai, Yutong He, Nicholas M. Boffi, Pradeep Ravikumar, Jose Miguel Hernandez-Lobato, Max Simchowitz, Benjamin Kurt Miller, Omar Chehab

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 날씨를 예측하려고 한다고 상상해 보세요. 당신에게는 내일 하늘이 어떤 모습일지 정확하게 알려줄 수 있는 아주 똑똑한 컴퓨터 모델(샘플)이 있습니다. 하지만 당신은 또한 그 특정 날씨 패턴이 실제로 일어날 확률이 얼마나 되는지(우도, likelihood)도 알아야 합니다.

인공지능의 세계에서, 고품질의 이미지나 분자 구조를 생성하는 것과 동시에 그것이 얼마나 발생할 확률이 높은지 정확하게 말할 수 있는 모델을 만드는 것은 매우 어려운 일입니다. 보통 우리는 둘 중 하나를 선택해야 합니다. 멋진 그림을 얻을 수는 있지만 확률을 계산할 수는 없거나, 확률을 계산할 수는 있지만 그림이 흐릿하거나 계산에 시간이 너무 오래 걸리는 방식 중 하나를 말이죠.

이 논문은 이 문제를 해결하기 위해 SCALLOP(Scalable Likelihood distillation of flOw maPs)이라 불리는 새로운 방법을 소개합니다. 이해를 돕기 위해 쉬운 비유를 들어 설명하겠습니다.

문제점: 수학의 "교통 체증"

강물이 잔잔한 호수(무작위 노이즈)에서 복잡한 도시(분자나 얼굴 같은 실제 데이터)로 흘러가는 모습을 상상해 보세요.

기존 방식: 흐르는 물의 양(우도)을 알기 위해, 과학자들은 강물을 멈추고 모든 방울을 측정하고 복잡한 계산을 수행해야 했습니다. 이것은 마치 교통 체증 속에서 모든 차를 멈춰 세워 차량 수를 세려는 것과 같았습니다. 이는 매우 느렸고, 종종 추측에 의존해야 했기에 엄청난 오차(높고 변동성이 큰 분산)를 낳았습니다.
"Hutchinson" 지름길: 이전의 시도들은 전체 교통량을 추측하기 위해 자동차의 일부를 무작위로 샘플링하여 속도를 높이려 했습니다. 하지만 이것은 거리의 단 세 명에게 물어보고 도시 전체의 자동차 수를 추측하는 것과 같았습니다. 때로는 작동하기도 했지만, 종-종 추측이 터무니없이 틀려 학습 과정을 불안정하고 노이즈가 심하게 만들었습니다.

해결책: SCALLOP의 "GPS와 속도계"

저자들은 엔진을 끄지 않고도 속도를 확인할 수 있는 GPS와 완벽한 속도계가 내장된 자동차와 같은 새로운 시스템을 구축했습니다.

플로우 맵 (GPS): 그들은 "플로우 맵"을 사용합니다. 이는 AI가 무작위 노이즈에서 실제 이미지나 분자로 이동하는 경로를 몇 단계 만에 알려주는 미리 계산된 지도와 같습니다 (마치 구불구불한 시골길 대신 고속도로를 타는 것과 같습니다).
우도 (속도계): 까다로운 부분은 동시에 "확률"(속도계의 눈금)을 아는 것입니다.
- 기존 방식: 기존 모델(F2D2 등)은 이 수치를 얻기 위해 앞서 언급한 "무작위 추측" 방식(Hutchinson 추정치)을 사용해야 했으며, 이는 노이즈가 심하고 느렸습니다.
- SCALLOP 방식: 저자들은 '조건부 발산 매칭(Conditional Divergence Matching)'이라는 수학적 트릭을 사용하여 추측 없이 직접 확률을 계산할 수 있다는 사실을 깨달았습니다. 자동차를 멈춰 세워 숫자를 세는 대신, 강물의 흐름 자체를 관찰하는 것입니다. 그들은 모델이 무거운 계산을 위해 멈출 필요 없이, 움직이는 과정 속에서 확률을 직접 "볼" 수 있게 해주는 공식을 도출했습니다.

이것이 왜 중요한가요?

논문은 SCALLлоP가 세 가지 측면에서 주요한 업그레이드라고 주장합니다.

더 부드럽습니다 (노이즈 감소): 무작위 추측에 의존하지 않기 때문에 학습 과정이 훨씬 안정적입니다. 속도계 바늘이 미친 듯이 요동치는 차를 운전하다가, 이제는 바늘이 일정하게 유지되는 차를 운전하는 것과 같습니다. 이는 모델이 더 빠르고 안정적으로 학습됨을 의미합니다.
더 빠릅니다: 기존 방식은 지도를 확인하기 위해 모든 빨간불마다 차를 멈춰야 하는 자동차와 같았습니다. SCALLOP은 직진합니다. 논문에 따르면 SCALLOP은 분자 과학 분야에서 기존의 최고 방법들보다 10배 더 빠르며, 다른 경쟁 모델들보다 100배 더 빠릅니다.
어디서나 작동합니다: 그들은 두 가지 매우 다른 대상에 대해 테스트했습니다.
1. 분자: 단백질(알라닌 펩타이드 등)이 어떻게 접히는지(folding) 시뮬레이션하는 데 사용했습니다. 이는 신약 개발에 매우 중요한데, 과학자들에게는 분자가 어떤 모양인지뿐만 아니라 그 모양으로 존재할 확률이 얼마인지도 중요하기 때문입니다. SCALLOP은 현재 기술 수준(state-of-the-art)보다 더 빠르고 정확하게 이를 수행했습니다.
2. 이미지: 얼굴을 생성하는 테스트(CelebA 데이터셋)도 진행했습니다. 고품질 이미지를 생성함과 동시에 정확한 확률 추정치를 보여줌으로써, 이 방법이 일반적인 데이터에도 적용 가능하다는 것을 증명했습니다.

결론

SCALLOP을 AI 생성기를 위한 새로운 엔진이라고 생각하세요. 이전에는 생성된 이미지나 분자의 "확률"을 알고 싶다면 시간과 정확도라는 막대한 세금을 지불해야 했습니다. SCALLOP은 이 세금을 없애줍니다. SC-ALLOP은 AI가 복잡한 데이터(3D 분자나 얼굴 등)를 생성하면서, 동시에 그것들이 존재할 확률이 얼마인지 한 번의 빠르고 안정적인 과정으로 알려줄 수 있게 해줍니다.

요약하자면: 이것은 AI가 무언가를 창조하고, 그 창조물이 존재할 확률이 얼마나 되는지 이해하도록 가르치는 더 빠르고, 더 정확하며, 노이즈가 적은 방법입니다.

기술 요약: 확장 가능한 라이클리후드 플로우 맵(SCALLOP)을 통한 소수 단계 볼츠만 생성기

문제 정의

최근 흐름 기반 생성 모델링(flow-based generative modeling)의 발전은 적은 수의 함수 평가(NFE)로도 고품질의 샘플 생성을 가능하게 했습니다. 그러나 모델의 우도(likelihood, 샘프트 밀도)를 추정하는 것은 여전히 병목 구간으로 남아 있습니다. 기존 방식들은 다음과 같은 트레이드오프에 직면해 있습니다:

제한적인 아키텍처: 정규화 흐름(Normalizing flows)은 변수 변환 공식을 통해 정확한 우도 계산을 보장하지만, 분자 모델링과 같이 대칭성이 있는 도메인에서 표현력을 제한하는 구조적 제약(예: 가역적 레이어)을 요구합니다.
고분산 추정기: 연속 정규화 흐름(CNF) 및 그 증류된 형태(예: Flow Maps)는 더 높은 표현력을 제공하지만, 속도장의 발산(divergence)을 적분해야 합니다. 이는 일반적으로 Hutchinson의 트레이스 추정기(trace estimator)를 사용하여 추정되는데, 이는 상당한 분산을 유발하며 차원이 높아질수록 계산 비용이 급격히 증가합니다.
증류의 한계: 최근의 F2D2(Ai et al., 2026)와 같은 방법들은 샘플링과 우도를 모두 몇 단계 내에 증류하려고 시도하지만, 여전히 훈련 과정에서 Hutchinson 기반의 발산 추정기에 의존하며, 이로 인해 발생하는 높은 분산과 계산 오버헤드를 그대로 물려받습니다.

결과적으로, 표현력이 뛰어난 소수 단계 생성 모델을 위한 확장 가능하고 저분산인 우도 추정 방법이 부족하며, 이는 자유 에너지 추정 및 향상된 샘플링과 같은 과학적 도메인에서의 볼츠만 생성기 활용을 저해하고 있습니다.

방법론: SCALLOP

저자들은 F2D2 라이클리후드 플로우 맵 모델을 기반으로 하되, Hutchinson 기반의 훈련 목적 함수를 확장 가능한 벡터화된 대안으로 대체한 SCAlable LikeLihood distillation of flOw maPs (SCALLOP) 프레임워크를 제안합니다.

1. 라이클리후드 플로우 맵 (Likelihood Flow Maps)

SCALLOP은 상태 $x_t$ (시간 $t$ 에서의 상태)를 상태 $x_s$ (시간 $s$ 에서의 상태)와 연관된 로그 밀도 변화량으로 공동 매핑하는 연산자인 라이클리후드 플로우 맵 $\Psi$ 를 사용합니다. 이는 다음과 같이 매개변수화됩니다:
$\Psi_\theta(x_t, t, s) = \begin{bmatrix} x_t \\ 0 \end{bmatrix} + (s-t) \begin{bmatrix} u_\theta(x_t, t, s) \\ D_\theta(x_t, t, s) \end{bmatrix}$
여기서 $u_\theta$ 는 속도(velocity)를 예측하고, $D_\theta$ 는 발산(divergence, 로그 밀도 변화)을 예측합니다.

2. 조건부 발산 매칭 (Conditional Divergence Matching, CDM)

핵률적인 혁신은 Hutchinson 트레이스 추정기를 **조건부 발산 매칭(CDM)**으로 교체하는 것입니다.

유도: 로그 밀도에 따른 흐름 궤적( $dt \log p_t(x_t) = \nabla \log p_t(x_t) \cdot v_t(x_t) + \partial_t \log p_t(x_t)$ )에 대해 연쇄 법칙(chain rule)을 적용하여 발산 항을 재구성합니다.
메커니즘: 선형 보간 설정( $x_t = (1-t)x_0 + tx_1$ )에서 타겟 데이터 포인트 $x_1$ 에 조건을 부여함으로써, 조건부 분포 $p_t(x_t|x_1)$ 를 다룰 수 있는 형태(가우시안)로 만듭니다. 이를 통해 명시적인 발산 연산자나 확률적 트레이스 추정기 없이도 발산을 표현할 수 있습니다.
목적 함수: CDM 목적 함수는 예측된 발산 $D_\theta$ 와 알려진 속도장 및 조건부 스코어로부터 유도된 항 사이의 차이를 최소화하며, 이를 통해 Hutchinson 추정기 자체를 우회합니다.

3. 벡터화된 발산 매칭 (Vectorized Divergence Matching, CDM-v)

효율성을 더욱 높이기 위해, 저자들은 **벡터화된 발산 매칭(CDM-v)**을 도입합니다.

네트워크는 스칼라 발산 $\nabla \cdot v_t$ 를 예측하는 대신, 야코비안(Jacobian)의 대각 성분인 $\text{diag}(\nabla v_t)$ 의 벡터를 예측합니다.
이를 통해 학습 대상이 스칼라에서 벡터 신호로 변환되어 데이터 효율성이 높아지며, 완전한 벡터화된 정식화가 가능해집니다.
이론적 분석에 따르면, 벡터화된 손실을 최소화하는 것은 상수 차이 내에서 표준 발산 매칭 손실을 최소화하는 것과 동일합니다.

4. 훈련 및 아키텍처

아키텍처: SCALLOP은 공유 인코딩 레이어와 샘플 속도 및 벡터화된 로그 밀도 변화를 위한 별도의 리드아웃 헤드(readout heads)를 갖춘 디퓨전 트랜스포머(Diffusion Transformer, DiT) 백본을 채택합니다.
대칭성: 분자 시스템의 경우, 아키텍처는 속도에 대해서는 SE(3) 등변성(equivariance)을, 로그 밀도 변화에 대해서는 (벡터 출력의 합산을 통해) SE(3) 불변성(invariance)을 강제합니다.
안정화: $t=1$ 에서의 특이점(singularity)을 처리하기 위해, $t$ 를 $[0, 1-\epsilon]$ 에서 샘플링하고 발산 매칭 항에 선형 가중치를 적용합니다.

주요 기여

확장 가능한 훈련 목적 함수: SCALLOP은 Hutchinson이 필요 없는 벡터화된 라이클리후드 증류 목적 함수(CDM-v)를 도입하여, F2D2에 비해 훈련 분산을 크게 줄이고 계산 비용을 절감했습니다.
경험적 우월성: SCALLOP은 유효 샘플 크기(Effective Sample Size), 에너지 Wasserstein-2, 비틀림 각도(Torsional Angle) Wasserstein-2 등 모든 지표에서 F2D2보다 우수한 성능을 입증했습니다. 또한 더 빠른 수렴 속도와 더 나은 테스트 세트 성능을 보여줍니다.
추론 효율성: 분자 도메인에서 SCALLOP은 정규화 흐름보다 최대 10배 빠르고, FALCON보다 100배 빠른 추론 성능을 달성하면서도 경쟁력 있거나 더 우수한 샘플 품질을 유지합니다.
일반화 능력: 고차원 이미지 데이터셋(CelebA-64)에 대한 검증을 통해, 정확한 라이클리후드 추정과 F2D2 대비 안정적인 훈련 역학을 보여주었습니다.

결과

분자 시스템 (볼츠만 생성기)

Alanine Dipeptide (ALA-2)부터 Hexapeptide (ALA-6) 시스템까지의 평가:

성능: SCALLOP은 모든 지표(Effective Sample Size, Energy Wasserstein-2, Torsional Angle Wasserstein-2)에서 F2D2를 일관되로 능가합니다. 또한 최첨단 모델인 FALCON 및 FALCON-A와 경쟁할 만한 성능을 달성했습니다.
안정성: SCALLOP의 훈련 그래디언트는 F2D2에 비해 현저히 낮은 분산과 높은 신호 대 잡음비(SNR)를 나타냅니다.
속도: SCLLOP은 추론 시 $10^4$ 의 큰 배치 크기를 지원할 수 있는 반면, FALCON은 야코비안 계산 비용으로 인해 매우 제한된 배치 크기를 요구합니다. SCALLOP은 정규화 흐름보다 10배 빠르고, FALCON보다 100배 빠릅니다.

이미지 생성 (CelebA-64)

SCALLOP은 1~8단계 오일러 스텝(Euler steps)에 대해 F2D2와 대등한 BPD(Bits Per Dimension) 및 FID를 달성했습니다.
F2D2와 비교하여 샘플당 BPD 오차를 일관되게 감소시켰습니다.
손실 계산은 16.9% 더 빠르며(F2D2 실행 시간의 83.1%), 더 안정적인 훈련 역학을 보여줍니다.

의의 및 주장

본 논문은 SCALLOP이 소수 단계 생성 모델에서 표현력과 다루기 쉬운 라이클리후드 평가 사이의 오랜 트레이드오프를 해결한다고 주장합니다. Hutchinson 트레이스 추정기를 제거함으로써, SCALLOP은 낮은 분산과 높은 확장성을 가진 매우 표현력 있는 라이클리후드 플로우 맵을 훈련할 수 있게 합니다.

저자들은 SCALLOP을 자유 에너지 계산 및 향상된 샘플링에 라이클리후드 인식이 필수적인 과학적 응용 분야를 위한 실질적인 솔루션으로 제시합니다. 저자들은 이론적으로 1단계로 샘플을 생성할 수 있지만, 고차원 작업에서는 실제적으로 약 10단계를 사용한다고 언급합니다. 이 방법은 속도장이 잘 학습되었다는 가정에 의존하며, 이 가정으로 인해 발생하는 편향(bias)에 대한 이론적 간극은 아직 정량화되지 않았습니다. 그럼에도 불구하고, 경험적 결과는 SCALLOP이 라이클리후드 인식이 가능한 생성 모델링을 가속화하는 경쟁력 있는 최첨단 접근 방식임을 입증합니다.