DiScoFormer: 분포 전반에 걸쳐 밀도(Density)와 스코어(Score)를 동시에 처리하는 단일 Transformer

📄 기술 보고서: arxiv.org/abs/2511.05924

머신러닝(Machine Learning)과 과학 분야의 많은 문제들은 결국 동일한 작업으로 귀결됩니다. 즉, 데이터 포인트의 집합을 가지고 그것들이 어디에서 왔는지, 즉 어떤 값이 흔하고 어떤 값이 드문지를 나타내는 분포(Distribution)를 복원하고자 하는 것입니다. 해당 분포를 확정한다는 것은 두 가지 양을 추정하는 것을 의미합니다. 바로 분포의 밀도(Density)와, 차원(Dimensionality)이 커질수록 더 유용하게 쓰이는 스코어(Score)입니다. 밀도는 히스토그램(Histogram)의 매끄러운 버전으로, 포인트들이 밀집된 곳에서는 높고 희소한 곳에서는 낮습니다. 로그 밀도(Log-density)의 기울기인 스코어는 밀도가 가장 빠르게 상승하는 방향을 가리킵니다. 스코어를 따라 포인트를 이동시키면 더 확률이 높은 영역으로 향하게 됩니다.

확산 기반 생성 모델(Diffusion-based generative models, Stable Diffusion이나 DALL-E와 같은 AI 이미지 생성기의 기반 기술)은 무작위 노이즈(Random noise)에서 시작하여 스코어를 반복적으로 따라가며, 그 노이즈를 현실적인 이미지로 변환합니다. 동일한 스코어는 베이지안 샘플링(Bayesian sampling)과 플라즈마(Plasma)와 같은 시스템을 모델링하는 데 사용되는 입자 시뮬레이션(Particle simulations)을 구동합니다.

유한한 샘플로부터 밀도와 스코어를 추출하는 것은 매우 어려운 과제이며, 오늘날의 도구들은 일반화 성능(Generalizability)과 정확도(Accuracy) 사이의 절충(Trade-off)을 강요합니다. 고전적인 접근 방식 중 하나인 커널 밀도 추정(Kernel Density Estimation, KDE)은 주변 데이터 포인트로부터 임의의 위치에서의 밀도를 계산합니다. 포인트들이 더 가깝고 많을수록 밀도는 높아집니다. 이 방식은 별도의 학습이 필요 없고 어떤 분포에도 적용할 수 있지만, 차원이 커짐에 따라 정확도가 급격히 떨어집니다. 대안으로, 스코어를 예측하도록 학습된 신경망 스코어 매칭(Neural score-matching) 모델은 고차원에서도 정확도를 유지하지만, 각 모델은 특정 분포를 학습해야 하며 다른 분포를 위해 처음부터 다시 학습해야 합니다.

우리는 DiScoFormer(Density and Score Transformer)라고 불리는 새로운 솔루션을 소개합니다. 이는 데이터 포인트 집합이 주어졌을 때, 재학습 없이 단 한 번의 순전파(Forward pass)만으로 분포의 밀도와 스코어를 모두 추정하는 단일 모델입니다.

DiScoFormer는 쌓인 Transformer 블록 레이어를 사용하여 전체 샘플을 그 배후 분포의 밀도(Density)와 스코어(Score)로 매핑합니다. 이 모델은 크로스 어텐션(Cross-attention)을 활용하여, 단순히 데이터가 존재하는 지점뿐만 아니라 임의의 지점에서도 밀도와 스코어를 평가할 수 있습니다. 스코어와 밀도는 수학적 관계를 공유하는데, 스코어는 밀도의 로그(Logarithm)에 대한 그래디언트(Gradient)입니다. 우리는 이를 활용하여 밀도를 위한 헤드와 스코어를 위한 헤드, 두 개의 출력 헤드를 가진 공유 백본(Shared backbone)을 사용합니다.

이러한 결합은 단순히 파라미터(Parameter)를 절약하는 것 이상의 역할을 합니다. 스코어 헤드는 모든 쿼리(Query)에서 로그 밀도 헤드의 그래디언트와 일치해야 하므로, 둘 사이의 어떠한 차이도 레이블이 필요 없는 일관성 손실(Label-free consistency loss)이 됩니다. 우리는 이를 추론(Inference) 시에 사용합니다. 컨텍스트(Context)를 고정하고 해당 일관성 손실에 대해 몇 번의 그래디언트 단계(Gradient steps)를 거치면, DiScoFormer는 정답 밀도(Ground-truth density)나 스코어가 필요하지 않은 상태에서도 즉석에서 분포 외(Out-of-distribution) 입력에 스스로 적응합니다.

Transformer 아키텍처가 이 작업에 적합한 데에는 수학적인 이유가 있습니다. 커널 밀도 추정(Kernel Density Estimation, KDE)은 단일 대역폭(Bandwidth)을 가집니다. 즉, 각 점의 영향력이 미치는 범위가 사전에 고정되어 모든 곳에 동일하게 적용됩니다. 어텐션(Attention)은 이를 엄격하게 일반화한 것입니다. 우리는 단일 어텐션 헤드의 가중치가 데이터에 대한 가우시안 커널(Gaussian kernel)과 거의 유사함을 분석적으로 보여주었으며, 따라서 하나의 크로스 어텐션 블록만으로도 이미 KDE의 밀도와 스코어를 재현할 수 있습니다. 여기서 모델은 더 나아가 여러 스케일(Scale)을 동시에 학습하고 이를 데이터에 적응시킵니다. DiScoFormer는 블랙박스를 위해 고전적인 방법을 버리는 것이 아니라, KDE를 특수한 사례로 포함하면서 이를 개선합니다.

DiScoFormer를 학습시키기 위해 어떤 데이터를 사용했을까요? 우리는 두 가지 주요한 이유로 가우시안 혼합 모델 (Gaussian Mixture Models, GMM)에 의존했습니다. 첫째, GMM은 보편적 밀도 근사기 (universal density approximators)입니다. 즉, 충분한 수의 성분 (components)이 있다면 임의의 작은 오차 범위 내에서 본질적으로 어떤 매끄러운 분포든 일치시킬 수 있습니다. 둘째, GMM은 폐형 밀도 (closed-form densities)와 스코어 (scores)를 가지고 있어, 우리가 감독 (supervise)할 수 있는 정확한 타겟을 항상 가질 수 있습니다. 우리는 매 배치(batch)마다 새로운 GMM을 추출함으로써 이 두 가지 속성을 모두 활용하며, 모델에 사실상 무제한의 타겟 분포 예시를 제공하고 각 예시를 주어진 GMM의 정확한 밀도 및 스코어와 대조하여 감독합니다.

전반적으로 DiScoFormer는 밀도 및 스코어 추정 모두에서 KDE를 능가하며, 특히 KDE가 어려움을 겪는 지점에서 그 격차는 더욱 벌어집니다. 100차원 환경에서는 차이가 압도적입니다. 가장 잘 튜닝된 KDE와 비교했을 때, DiScoFormer는 스코어 오차를 약 6.5배 줄이고 밀도 오차를 37배 이상 줄였으며, KDE가 메모리 부족 문제를 겪는 동안에도 샘플을 추가할수록 성능이 계속 향상됩니다. 또한 학습 데이터의 범위를 훨씬 벗어나서도 작동하며, 학습 중에 보았던 것보다 더 많은 모드 (modes)를 가진 혼합 분포나 라플라스 (Laplace) 및 스튜던트 t (Student-t) 분포와 같은 비가우시안 (non-Gaussian) 형태에서도 정확도를 유지합니다. KDE의 주요 장점은 여전히 속도이며, 특히 데이터셋이 작을 때 그러합니다.

우리가 DiScoFormer에서 가장 유망하다고 생각하는 부분은 스코어 추정 (score estimation)이 생성 모델링 (generative modeling), 베이지안 추론 (Bayesian inference), 과학 계산 (scientific computing) 등 많은 분야에서 공유되는 의존성이라는 점입니다. 고차원에서도 정확도를 유지하며 문제마다 다시 학습할 필요를 없애주는 사전 학습된 플러그인 추정기 (pretrained, plug-in estimator)는 이 모든 분야의 비용을 한 번에 절감할 수 있을 것입니다. 즉, 스코어와 밀도가 나타나는 모든 곳에서 재사용 가능한 단 하나의 모델이 되는 것입니다.

더 자세한 내용은 저희의 기술 보고서 (technical report)를 읽어보시길 권장합니다.