디스코포머 변환기로 구현하는 밀도와 스코어 추정 플러그인

초록

DiScoFormer는 한 번 학습한 트랜스포머 모델을 이용해 임의의 i.i.d. 샘플 집합으로부터 확률밀도와 그 스코어(∇log f)를 동시에 추정한다. 모델은 샘플 순열에 대해 equivariant하고, 입력을 정규화해 affine equivariance를 만족한다. 이론적으로 셀프‑어텐션이 정규화된 커널 밀도 추정(KDE)을 재현함을 증명하고, 실험에서는 다중 스케일 헤드가 커널‑유사 행동을 학습함을 확인한다. 결과적으로 DiScoFormer는 기존 KDE보다 빠르게 수렴하고, 고차원·대규모 데이터에서도 정확한 밀도·스코어 추정을 제공한다. 또한 추정된 스코어를 이용해 스코어‑디바이즈드 KDE, 피셔 정보, Fokker‑Planck 방정식 등 다양한 downstream 작업에 바로 활용할 수 있다.

상세 분석

본 논문은 확률밀도와 그 로그 미분(스코어) 추정을 하나의 연산자로 정의하고, 이를 “시퀀스‑투‑오퍼레이터” 학습 문제로 전환한다. 핵심 아이디어는 i.i.d. 샘플 집합 X = {x₁,…,xₙ}을 트랜스포머의 입력 시퀀스로 보고, 동일한 네트워크가 전체 집합에 대해 f(xᵢ)와 ∇log f(xᵢ)를 동시에 출력하도록 학습하는 것이다. 이를 위해 두 가지 대칭성을 강제한다. 첫째, 샘플 순열에 대한 equivariance는 트랜스포머의 순열 불변 구조와 positional encoding을 배제함으로써 자연스럽게 만족된다. 둘째, affine equivariance는 입력을 화이트닝(공분산 정규화)하고, 학습 과정에서 무작위 회전·스케일 변환을 적용함으로써 구현한다. 논문은 Proposition 3.1을 통해 이러한 대칭성이 수학적으로 보장됨을 증명한다.

이론적 기여 중 가장 눈에 띄는 부분은 셀프‑어텐션이 정규화된 Gaussian KDE와 동일한 가중치를 생성한다는 증명이다. Proposition 3.2는 Q = K = h·X, V = Iₙ 형태의 단일 어텐션 헤드가 Softmax(QKᵀ)·V를 통해 각 샘플 간 거리 기반 가우시안 커널을 정확히 구현함을 보여준다. 따라서 트랜스포머는 기존 KDE의 비선형 스무딩을 넘어, 다중 헤드와 다층 구조를 통해 다중 스케일·다중 방향 커널을 학습적으로 생성할 수 있다. 실험적인 시각화(Figure 1, 2)에서는 일부 헤드가 장거리, 일부 헤드가 근거리, 또 다른 헤드가 특정 방향에 집중하는 현상이 관찰되어, “헤드 특화”가 자연스럽게 발생함을 확인한다.

학습 데이터는 폐쇄형 형태의 Gaussian Mixture Model(GMM)에서 샘플을 동적으로 생성한다. GMM은 연속적으로 미분 가능한 밀도와 스코어를 제공하므로, 모델이 정확한 라벨을 얻을 수 있다. 손실 함수는 밀도와 스코어 MSE를 가중합한 형태이며, α 파라미터를 통해 두 목표 간 균형을 조절한다. 모델 아키텍처는 4개의 트랜스포머 인코더 레이어(히든 차원 128, 8 헤드)와 두 개의 출력 헤드(밀도, 스코어)로 구성되며, 파라미터 수는 약 80만 개에 불과하다.

실험에서는 1차원부터 10차원까지, 샘플 수 256부터 8192까지 다양한 설정에서 기존 KDE와 Score‑Debiased KDE(SD‑KDE)를 능가하는 성능을 보였다. 특히 고차원(10D)에서 KDE는 심각한 차원 저주 현상을 보이지만, DiScoFormer는 상대적으로 낮은 MSE를 유지한다. 또한 훈련 시 1–10 모드 GMM만 사용했음에도 불구하고, 1–19 모드의 테스트 GMM에 대해 안정적인 일반화 능력을 보여준다. 추정된 스코어는 SD‑KDE에 직접 플러그인될 수 있어, 기존 방법보다 더 정확한 밀도 추정이 가능하고, 피셔 정보 계산이나 Fokker‑Planck 방정식의 수치 해석에도 바로 활용될 수 있다.

전체적으로 이 논문은 (1) 트랜스포머가 비파라메트릭 밀도·스코어 추정에 적합한 구조임을 이론·실험적으로 입증하고, (2) 셀프‑어텐션이 커널 기반 방법의 일반화된 형태임을 보여주며, (3) “train‑once, infer‑anywhere”라는 실용적인 프레임워크를 제공한다는 점에서 의미가 크다. 향후 연구는 더 복잡한 비정규분포, 비정상적 데이터(예: 이미지, 시계열)로의 확장과, 연산 효율성을 위한 희소 어텐션 또는 저차원 임베딩 기법과의 결합이 기대된다.