기술 요약: 오프-폴리시 로그 분산 정규화를 이용한 볼츠만 생성기(Boltzmann Generators)의 효율적 학습

1. 문제 정의

비정규화된 확률 밀도, 특히 물리계의 볼츠만 분포로부터 샘플링하는 것은 계산 물리학 및 화학의 근본적인 과제이다. 대상 분포는 $p_X(x) = \tilde{p}_X(x)/Z$로 정의되며, 여기서 $\tilde{p}_X(x) = \exp(-E(x)/k_B T)$는 비정규화된 밀도이고, $E(x)$는 에너지 함수이며, $Z$는 다루기 어려운 정규화 상수이다.

분자 역학(MD) 및 마르코프 연쇄 몬테카를로(MCMC)와 같은 고전적인 방법들은 점근적으로는 정확하지만, 특히 고차원이거나 에너지가 복잡한 시스템에서는 긴 상관 관계가 있는 궤적으로 인해 느린 혼합(mixing)과 높은 계산 비용을 겪는 경우가 많다. 볼츠만 생성기(Boltzmann Generators, BGs)는 생성 모델로서의 대안을 제공하여 평형 분포로부터 독립적인 샘플링을 가능하게 한다. 그러나 이들의 실질적인 유용성은 데이터 효율성에 의해 제약을 받는다. 표준적인 학습 목적 함수는 일반적으로 다음 사항들에 의존한다:

1. 대상 분포로부터 얻은 평형 샘플 (MD를 통해 생성하기에는 비용이 많이 드는 경우가 많음).
2. 변분 설정(variational settings)에서의 광범위한 에너지 평가.

기존의 접근 방식들은 데이터셋 생성 중에 사용 가능한 대상 에너지 라벨을 충분히 활용하지 못하거나, 이 정보를 효과적으로 통합하기 위해 추가적인 온-폴리시(on-policy) 샘플과 에너지 평가를 요구하는 경우가 많다.

2. 방법론: 오프-폴리시 로그 분산 정규화 (LDR)

저자들은 **오프-폴리시 로그 분산 정규화(Off-Policy Log-Dispersion Regularization, LDR)**를 제안한다. 이는 추가적인 온-폴리시 샘플이나 추가적인 에너지 평가 없이도 대상 에너지 정보를 볼츠만 생성기 학습에 통합하는 프레임워크이다.

핵심 개념

이 방법은 온-폴리시 변분 학습에서 전통적으로 사용되는 로그 분산(log-variance) 목적 함수를 기반으로 한다. 저자들은 $-\log q_\theta^X(x)$와 $-\log \tilde{p}_X(x)$의 차이를 나타내는 함수 $f_\theta(x)$를 다음과 같이 정의한다 (비정규화된 중요도 가중치):
$$f_\theta(x) = -\log q_\theta^X(x) - (-\log \tilde{p}_X(x)) = -\log q_\theta^X(x) - \frac{E(x)}{k_B T}$$
최적점($q_\theta^X = p_X$인 지점)에서 $f_\theta(x)$는 상수($\log Z$)가 되어야 한다.

로그 분산 목적 함수

저자들은 로그 분산 목적 함수를 다음과 같이 로그 분산 목적 함수의 가족으로 일반화한다:
$$\mathcal{L}_\theta^{LD} = \mathbb{E}_{r_X} \left[ \rho(f_\theta(x) - c_\theta) \right], \quad c_\theta = \mathbb{E}_{r_X}[f_\theta(x)]$$
여기서 $\rho(u) \geq 0$는 분산 척도(예: $|u|^p$)이며, $r_X$는 참조 분포이다.

• LDR-L2 ($p=2$): 분산(제2 중심 모멘트)을 최소화하며, 이는 표준 로그 분산 목적 함수를 회복한다.
• LDR-L1 ($p=1$): 제1 절대 중심 모멘트를 최소화한다. 저자들은 이 변형이 에너지 아웃라이어(outlier)에 덜 민감하며, 고차원 에너지 지형의 헤비 테일(heavy-tailed) 특성을 고려할 때 더 안정적인 최적화를 제공한다고 가설을 세웠다.

오프-폴리시 정규화 전략

참조 분포 $r_X$(예: 고정된 데이터셋)가 구성 공간 전체에 대한 풀 서포트(full support)를 갖지 않는 경우, 로그 분산 목적 함수만을 최소화하는 것은 발산(divergence)이 아니라는 중요한 이론적 통찰(부록 A)이 있다. 즉, 데이터 매니폴드 외부의 모델을 제약하지 못하여 잘못된 정규화를 초래한다.

이를 해결하기 위해, LDR은 표준 데이터 기반 목적 함수(예: Forward KL, Flow Matching, 또는 Score Matching)와 함께 정규화 항으로 적용된다:
$$\mathcal{L}_\theta = \lambda_{data}\mathcal{L}_\theta^{data} + \lambda_{LD}\mathcal{L}_\theta^{LD}$$

• 데이터 기반 항은 전역적인 정규화와 대상 분포로의 수렴을 보장한다.
• LDR 항은 데이터 매니폴드 상에서 로그 중요도 가중치의 분산을 최소화함으로써 제안 분포(proposal distribution)를 대상 에너지 지형과 정렬시키는 형상 정규화(shape regularizer) 역할을 한다.

이 접근 방식은 추가적인 타겟 에너지 평가 없이도 편향된 시뮬레이션 데이터나 보조 에너지 라벨 데이터셋을 포함한 임의의 참조 분포를 사용할 수 있게 해준다.

3. 주요 기여

1. 오프-폴리시 설정으로의 적응: 저자들은 고정된 에너지 라벨 데이터셋 상에서 로그 분산 목적 함수를 온-폴리시 변분 학습에서 오프-폴리시 설정으로 적응시켜, 이를 정규화제로 활용하였다.
2. 로그 분산으로의 일반화: 로그 분산 손실 함수를 하나의 가족으로 일반화하여, 표준 L2 변형보다 아웃라이어에 대해 더 나은 안정성을 제공하는 L1 변형(LDR-L1)을 도입하였다.
3. 폭넓은 적용성: 이 프레임워크는 다음 세 가지 뚜렷한 학습 레지임에서 효과적임을 입증하였다:
   • 비편향 평형 데이터에 대한 학습.
   • 편향된 시뮬레이션 데이터(중요도 샘플링 및 혼합 데이터셋 활용)에 대한 학습.
   • 대상 샘플에 접근할 수 없는 순수 변분 학습(대상 에너지 평가만을 사용).
4. 최신 기술과의 통합: 이 방법은 현재 변분 학습의 최첨단 기술인 제약된 질량 수송(Constrained Mass Transport, CMT)과 성공적으로 통합되어 효율성을 크게 향상시켰다.

4. 실험 결과

저자들은 2D 가우시안 혼합 모델(GMM), 알라닌 디펩티드($d=60$), 알라닌 헥사펩티드($d=180$)를 포함한 여러 벤치마크에서 LDR을 평가하였다. 성능은 음의 로그 가능도(NLL)와 유효 샘플 크기(ESS)를 사용하여 측정되었다.

주요 결과:

• 비편향 데이터: 비편향 데이터셋에서 LDR은 NLL과 ESS를 모두 크게 개선하였다. 알라닌 디펩티드의 경우, LDR은 $1 \times 10^6$개의 샘플만으로도 정규화되지 않은 학습이 $5 \times 10^6$개의 샘플을 사용했을 때의 성능을 능가하였으며, 이는 5배의 데이터 효율성 향상을 의미한다.
• 편향된 데이터: 짧은 궤적을 통해 생성된 편향된 데이터셋으로 학습할 때, 재표집된 평형 데이터셋($D_{IS}$)과 원래의 편향된 데이터셋 모두에 LDR을 적용했을 때 상당한 이득이 있었다. LDR은 $1 \times 10^5$개의 중요도 샘플과 초기 편향 데이터셋 구성을 사용하여 고성능 모델을 학습할 수 있게 해주었으나, 정규화되지 않은 베이스라인은 이보다 10배 이상 많은 중요도 샘플을 필요로 했다.
• 변분 학습 (대상 샘플 없음): CMT를 사용하는 순수 변분 설정에서, LDR은 필요한 대상 에너지 평가 횟수를 약 한 자릿수(one order of magnitude) 줄였다. 알라닌 디펩티드의 경우, CMT + LDR은 표준 CMT가 $1 \times 10^8$번의 평가를 사용하는 대신 $1 \times 10^7$번의 평가만으로도 대등한 성능을 달 achieved 하였다. 알라닌 헥사펩티드의 경우, LDR은 바닐라 CMT가 실패한 레지임(ESS < 6%)에서도 모드 붕괴를 방지하며 우수한 성능(ESS > 30%)을 보였다.
• 데카르트 좌표: TarFlow를 사용하여 데카르트 좌표에서 직접 학습할 때도 이 이득이 재현되었으며, 이는 혜택이 내부 좌표 표현에 국한되지 않음을 확인해 준다.
• 안정성: LDR-L1은 특히 헤비 테일 특성을 가진 고차원 시스템에서 LDR-L2보다 더 안정적인 학습 역학(더 낮은 그래디언트 노름 분산)을 제공하였다.

5. 의의 및 주장

본 논문은 LDR이 볼츠만 생성기 학습에 대상 에너지 정보를 통합하는 단순하면서도 강력한 프레임워크를 제공한다고 주장한다. 이 연구의 주요 의의는 다음과 같다:

• 데이터 효율성: LDR은 다양한 설정에서 최종 성능과 데이터 효율성을 일관되게 개선하며, 종종 필요한 샘플 수나 에너지 평가 횟수를 한 자릿수 수준으로 줄여준다.
• 유연성: 비편향, 편향 및 순수 변분 학습을 모두 지원하며, 추가적인 온-폴리시 샘플이나 추가적인 에너지 평가 없이도 광범위하게 적용 가능하다.
• 상보성: 저자들은 LDR이 경로 그래디언트(path gradients, 타겟 그래디언트를 사용)와 같은 기존의 분산 감소 기법과 상보적이라고 언급한다. LDR은 타겟 에너지 값(zero-order 정보)을 활용하기 때문이다.

저자들은 LDR이 데이터 크기에 대한 의존성을 제거하지는 못하지만, 각 샘플로부터 추출하는 정보를 극대화함으로써 거대한 데이터셋의 필요성을 크게 완화한다는 점을 강조한다. 이 프레임워크는 기존의 노멀라이징 플로우(normalizing flows) 및 잠재적으로 확산 기반 생성기(diffusion-based generators)의 학습 파이프라인에 쉽게 통합될 수 있도록 설계되었다.