[평범한 학부생이 하는 논문 리뷰] A Geometric View of Data Complexity : Efficient Local Intrinsic Dimension Estimation with Diffusion Models (NeurIPS 2024)

Paper : https://arxiv.org/abs/2406.03537

A Geometric View of Data Complexity: Efficient Local Intrinsic Dimension Estimation with Diffusion Models

Abstract

 LID estimation을 위해 deep generative model을 이용한다. 하지만 generative model을 이용한 최근 method들은 부정확한 estimate을 제공하거나, single pre-trained model 이상을 필요로 하거나, computationally intensive하거나, diffusion model과 같은 best available deep generative models를 이용하지 않는다. 본 논문에서는 DM과 연관된 Fokker-Planck equation이 앞선 문제를 다루는 LID estimator를 제공할 수 있다는 것을 보인다.

1. Introduction

 Manifold hypothesis는 $\mathbb{R}^D$에서 high-dimensional data는 $\mathbb{R}^D$의 low-dimensional submanifolds에 종종 존재한다라고 말한다. 이 가설은 데이터의 복잡성의 natural measure로 local intrinsic dimension(LID)을 사용하도록 한다. $\text{LID}(x)$는 $x$가 속하는 data manifold의 dimension에 대응된다. 그리고 이는 $x$를 설명하는데 필요로 되는 최소한의 variable의 수로 이해될 수 있다. Data manifold는 명시적으로 알 수 없어 LID는 추정되야만 한다. 본 논문에서는 'query datum $x$가 따르는 dataset이 주어질 때, 어떻게 tractably하게 $\text{LID}(x)$를 추정할 것인가?'라는 문제를 다룬다.

 기존의 model-free 방식은 large dataset에 대해서는 계산하는데 비싸다는 문제를 가졌고, 기존의 model-based 방식 또한 부정확하다, computationally expensive, 존재하는 best generative model을 사용하지 않는다, 여러 모델들을 학습할 필요가 있다, pre-trained model에 의존하기보다는 training 절차를 바꾸는 등의 문제점을 가지고 있다. 게다가, 이러한 방식들 중에 어떤 것도 Stable Diffusion에 의해 생성된 이미지들처럼 고해상도의 이미지로 확장되지 못했다.

 본 논문은 model-based estimator인 LIDL을 기반으로, 하나의 pre-trained diffusion model(DM)만을 사용해서 LID가 어떻게 효율적으로 추정되는지를 보여줌으로써 이러한 문제들을 해결한다. LIDL data에 다른 수준의 Gaussian noise를 곱하고, 각 수준에 대해서 normalizing flow를 학습시키고, noise의 log standard deviation를 covariate으로 하고 이에 대응하는 response로 $x$에서 평가된 log density를 사용해서 linear regression을 fitting하는 것으로 계산된다. 이를 통해 얻어진 slope는 $\text{LID}(x) - D$의 추정치가 된다. 본 논문에서는 먼저 많은 normalizing flow보다 하나의 DM만이 필요한 방식으로 어떻게 LIDL을 DM에 적용할지 보여준다. 이를 직접 적용하는 것은 하나의 DM만 필요로 하지만 여러번의 ODE solver의 호출을 필요로 한다. 본 논문은 한번의 solve로 모든 필요로 되는 log density를 계산하는 대안의 ODE로 이를 어떻게 우회할지 보인다. 그런 다음 본 논문은 LIDL에서 regression의 slope이 diffusion process에서 marginal log probability의 변화 비율을 포착하는 것을 목표로 한다는 것을 논의한다. 이러한 결과로 나온 estimator FLIPD는 매우 효율적이고 ODE solver에 대한 필요를 우회한다.

Contributions

1. 어떻게 DM이 ODE solver에 대한 한번의 호출로 LIDL과 효과적으로 결합되는지 보임
2. FLIPD를 제안하기 위해서 Fokker-Planck equation을 활용하고, estimator를 향상시키고 ODE solver에 대한 필요를 위회함
3. 이론적으로 FLIPD를 동기부여함
4. 이전의 evaluation에서의 gap을 드러내고 특히 manifold의 복잡성이 증가함에 따라 다른 estimator에 대해서는 정확하지 않은 LID estimation benchmark task의 확장된 suite을 도입
5. DM에 대해서 fully-connected architecture를 사용할 때, FLIPD는 더 computationally efficient하면서 기존의 baseline의 성능을 능가함
6. Natural image들에 적용될 때, FLIPD estimates는 다른 complexity measure와 align
7. SD의 latent space에 적용될 때, FLIPD는 처음으로 고해상도의 image에 대해서 LID를 추정 가능

2. Background and Related Work

2.1 Diffusion Models

Density Evaluation

DMs는 continuous normalizing flows로 해석될 수 있고, 그래서 density evaluation이 허용된다. 이는 $\hat{Y}_{1-t}$의 분포 $\hat{p}(\cdot, t)$이라고 하면, $\hat{p}(x, t)$는 어떤 $x \in \mathbb{R}^D, t_0 \in (0,1]$에 대해서도 수학적으로 evaluate될 수 있다. 이는 DM과 연관된 (forward)ODE 덕분에 달성된다.

$t_0$에서 1까지 이 ODE를 푸는 것은 trajectory $(\hat{x}_t)_{t \in [t_0,1]}$를 생성하고 이는 continuous change-of-variable formula를 통해 density evaluation에 사용된다.

$v(x,t) := f(x,t) - g^2(t)\hat{s}(x,t)/2$

Trace estimation

특정 $\hat{x}_t$에 대한 $\nabla v(\hat{x}_t,t)$를 계산하는 것에 대한 비용은 $\Theta(D)$이다. (5)에서 integral을 계산하기 위해서, $(\hat{x}_t)_{t \in [t_0,1]}$는 길이 $N$의 trajectory로 discretize돼야한다. $v(\hat{x}_t,t)$를 evaluate하는데 드는 비용이 $F$라고 하면, deterministic density evaluation은 $\Theta(NDF)$이다. $tr(M) = \mathbb{E}_\epsilon[\epsilon^TM\epsilon]$을 사용하는 Hutchinson trace estimator는 그래서 stochastic density estimation에서 흔히 사용된다. 이는 $\Theta(NkF)$이고, $k \ll D$이다.

2.1 Local Intrinsic Dimension and How to Estimate it

LID

Disjoint union of manifold와 이 union에 속하는 point $x$가 주어졌을 때, $x$의 local intrinsic dimension는 $x$가 속한 submanifold의 dimension이다. LID는 point $x$의 intrinsic property가 아니라, $x$를 포함한 manifold에 대한 $x$의 property이다. $LID(x)$는 $x$를 포함하는 manifold에 존재하는 variation의 요소들의 수와 대응된다. 그래서 $x$의 상대적인 복잡성에 대한 natural measure다.

Estimating LID

Generative model들은 $p(\cdot, 0)$를 학습하기 때문에 이들은 LID를 측정하기 위한 기존 방식들의 대안이 된다. 학습이 잘 됐다면, 그들은 $p(\cdot,t)$의 support에 대한 정보를 encode해야한다. Generative model들로부터 이러한 정보를 뽑는 것은 쉽지 않다. 

LIDL

LIDL은 tractable density estimators로써 normalizing flows에 의존한 LID estimation method이다. 본 논문에서는 $p(\cdot,0)$과 log standard deviation $\delta$인 Gaussian noise의 convolution을 다음과 같이 나타낸다.

$p(\cdot,0)$에 대한 mild regularity condition하에서, 그리고 $p(\cdot,0)$의 support안의 $x$가 주어졌을 때, $\delta \rightarrow - \infty$에 대해서 다음을 따른다.

$\delta$의 negative enough value에 대해서 다음과 같다.

만약 다양한 $\delta$ 값에 대해서 $\text{log} \varrho(x,\delta)$를 evaluate할 수 있다면, $LID(x)$를 추정하는데 avenue를 제공할 것이다. $\delta_1, \cdots, \delta_m$의 값을 셋팅하고, $\{ (\delta_i, \text{log} \varrho(x,\delta_i)) \}^m_{i=1}$로 linear regression을 fitting한다. 이 regression의 slope를 $\hat{\beta}_x$라고 하면, 이는 $LID(x) - D$를 추정하고, 그래서 $LID(x) \approx D + \hat{\beta}_x$가 local intrinsic dimension에 대한 sensible estimator가 된다.

$\varrho(x,\delta)$는 unknown이고 evaluate할 수 없기 때문에, $m$개의 normalizing flow를 학습할 필요가 있다. 이는 LIDL의 단점으로 꼽히고 있다. 또한 low-dimensional manifold를 학습하는데 어려움을 겪는다고 한다. 이러한 문제는 DM을 통해 해결될 수 있다고 한다.

Estimating LID with DMs

$t \searrow 0$에 따라, score function $s(x,t)$는 $x$를 포함하는 manifold를 orthogonal하게 지시한다. 그래서 기존 연구는 normal bundle(NB) estimator라는 LID estimator를 제안한다.

$x$에서 시작해서 $t_0$까지 (1)를 실행한다. 최종 값에서 $\hat{s}(\cdot,t)$를 계산한다. 이를 $K$번 반복해 얻은 vector들을 stack해서 $S(x) \in \mathbb{R}^{D \times K}$를 얻는다. $t_0$이 충분히 작고 $K$가 충분히 크면, $S(x)$의 column들은 $x$에서 manifold의 normal space를 span한다. 이는 이 행렬의 rank가 $D - LID(x)$인 이 normal space의 dimension을 추정한다는 것을 제안한다.

수치적으로, rank는 $S(x)$의 SVD를 수행하고 threshold를 넘는 singular value의 수를 센다. 이 NB estimator는 $\Theta(KF + D^2K)$만큼의 비용이 든다. 이 방식은 기존의 LIDL의 몇몇 한계점을 해결하지만, high dimension에서 computationally expensive하다.

3. Method

기존 논문에서는 LIDL에 normalizing flows만을 사용했지만, 이를 DMs으로 대체할 수 있다. 본 섹션을 통해 본 논문에서는 LIDL에 DMs을 naive하게 적용하는 것에 progressive improvements를 보여준다. 본 논문에서는 drift term $f(x,t) = b(t)x$를 가정하는 pre-trained DM을 사용한다. 이러한 선택은 (1)과 연관된 transition kernel $p_{t|0}$이 Gaussian이라는 것을 암시한다.

여기서 $\psi, \sigma$는 미분 가능하고 $\lambda(t) := \sigma(t)/\psi(t)$로 단사함수(서로 다른 입력값에 대해 다른 출력값을 가지는 함수)이다. 이러한 셋팅은 실제로 흔히 사용되는 모든 DMs를 통합한다.

3.1 LIDL with a Single Diffusion Model

Single DMs는 이미 다양한 수준의 noise와 data를 convolve함으로써 작동하고 이는 (5)와 같이 resulting noisy distribution의 density evaluation을 가능하게 한다. 이러한 이유로, 본 논문은 single DM을 이용해서 LIDL이 가능하다는 것을 보인다. 필요한 것은 $\varrho(\cdot,\delta)$를 DM의 density $p(\cdot,t)$와 연관 짓는 것이다. $\psi(t) = 1$인 variance-exploding DMs의 경우에서, (10)의 transition kernel의 defining property를 쉽게 사용해 다음 수식을 얻을 수 있다.

이는 $t = \sigma^{-1}(e^\delta)$라 설정하면 (6)에서의 $\varrho(x,\delta)$와 동일하다. 결국, (5)를 통해 각 $\hat{p}(x,\sigma^{-1}(e^{\delta_i}))$를 evaluate함으로써 single variance-exploding DM으로 LIDL을 사용할 수 있다. 본 논문에서는 (10)과 같은 transition kernel을 가진 임의의 DM에 대해서도, 다음을 보인다.

여기서 $t(\delta) := \lambda^{-1}(e^\delta)$이다. LIDL은 $\text{log} \varrho (\cdot, \delta)$를 필요로 하고, DMs는 $\text{log} p(\cdot,t)$를 제공한다. 이 두 quantity를 연결함으로써 LIDL이 single DM를 이용해 사용될 수 있다는 것을 보이기 때문에, 위 방정식은 중요한 의미를 갖는다.

3.2 A Better Implementation of LIDL with a Single Diffusion Model

LIDL과 함께 (12)를 사용하는 것은 regression을 하기 전에 각 $i = 1, \cdots,m$에 대해서 (5)를 통해 $\text{log} \hat{p}(\psi(t)(\delta_i))x,t(\delta_i))$를 계산하는 것을 여전히 포함한다. (4)에서 각각의 대응되는 ODEs는 다른 시간 $t_0 = t(\delta_i)$에서 시작하고 다른 지점 $\psi(t(\delta_i))x$에서 evaluate되기 때문에, 각 $i$에 대해서 다른 ODE solver 호출이 사용되야하는 것을 의미하고 이는 과도하게 expensive procedure를 야기한다. 이를 해결하기 위해서 본 논문에서는 $\partial/\partial \delta \ \text{log} \varrho(x,\delta)$에 대한 explicit formula를 찾는 것을 목표로 한다. 본 논문에서는 $\partial/\partial \delta \ \text{log} p(x,t)$에대한 명시적인 formula를 제공하는 (1)과 연관된 Fokker-Planck equation을 활용한다. Chain rule과 (12)와 함께 이 방정식을 사용해, (10)에서처럼 transition kernel을 가진 DMs에 대해 다음을 보인다.

$\delta_1 < \cdots < \delta_m$을 가정하면, 위의 수식을 사용해서 ODE를 통해 $\text{log} \hat{\varrho}(x,\delta)$를 다음과 같이 정의할 수 있다.

$\delta_1$에서 $\delta_m$로 이 ODE를 푸는 것은 trajectory $(\text{log}\hat{\varrho}(x,\delta))_{\delta \in [\delta_1, \delta_m]}$를 생성한다.  $v(t(\delta);\hat{s},x)$가 $\hat{\varrho}(x, \delta)$에 의존하지 않기 때문에, $\hat{s} = s$일 때 위 ODE의 solution은 상수 차이만 가진다. $$\text{log} \hat{\varrho}(x,\delta) =  \text{log} \varrho(x,\delta) + c_{init}$$

몇 $c_{init}$은 ODE의 초기 조건에는 의존하지만 $\delta$에 depend하지 않는다. 게다가 LIDL은 $\{ (\delta_i, \text{log} \hat{\varrho} (x,\delta_i)) \}^m_{i = 1}$을 사용해서 regression을 fit하는데, 그래서 $c_{init}$은 slope에 영향을 끼치지 않고 intercept에 흡수되기 때문에, initial condition을 0으로 설정하는게 이상해 보일 수 있지만 그렇지 않다. 즉, initial condition은 무관하고 (14)에 대한 ODE solver만 한번 호출하면 DMs로 LIDL을 사용할 수 있다.

3.3 FLIPD : An Efficient Fokker-Planck-Based LID Estimator

(14)를 푸는 것은 하나의 ODE solver에서 $\hat{s}$의 Jacobian의 trace를 여러 번 계산하는 것을 포함하고 이는 여전히 비싸다. 본 논문에서는 ODE solver에 대한 필요를 우회하는 LDI estimator FLIPD를 제안한다. LIDL은 (8)을 기반으로 하는데, 이 방정식을 미분하는 것은 negative enough $\delta_0$에 대해서 $\partial / \partial \delta \ \text{log} \varrho(x,\delta_0) \approx LID(x) - D$를 생성한다. 이는 (13)가 LIDL에서 regression이 추정하려고 하는 변화의 비율을 직접적으로 제공한다는 것을 의미한다.

여기서 $t_0 := t(\delta_0)$이다. $\hat{s}$의 Jacobian의 trace는 $v(t(\delta_0);\hat{s},x)$를 계산할 때 한번만 계산되기 때문에 FLIPD를 계산하는 것은 매우 싸다. 앞서 2.1에서 stochastic Hutchinson trace estimator를 사용하면 cost를 $\Theta(DF)$를 $\Theta(kF)$로 줄일 수 있다. 특히 NB estimator의 $\Theta (KF + D^2K)$에 비해 훨씬 빠른 속도이다. ODE solver를 필요로하지 않을 뿐만 아니라, FLIPD를 계산하는 것은 single hyperparameter $\delta_0$ 하나만 셋팅하는 것을 필요로 한다. 게다가, $v(t(\delta); \hat{s},x)$는 $t(\delta)$통해서만 $\delta$에 의존하기 때문에, $\delta_0$ 대신에 $t_0$를 hyperparameter로 직접 셋팅할 수 있다. 이는 $t(\delta_0) = \lambda ^{-1}(e^{\delta_0})$에 대한 복잡한 계산을 피할 수 있다. $\rightarrow$ 적절하게 negative $\delta_0$를 셋팅하는 대신에, 0에 충분히 가까운 $t_0 > 0$를 셋팅한다.

본 논문은 FLIPD의 이론적인 정당성을 제시한다. (7)에서 $\mathcal{O}(1)$ term이 (8)에서와 같이 $\delta$에 대해 상수일 필요는 없다. 이 term이 constant로부터 멀수록 LIDL과 FLIPD 모두에 대해서 bias가 커진다. 다음 결과는 idealized linear setting에서 FLIPD는 이 문제에 영향을 받지 않는다는 것을 보인다.

본 논문은 이 theorem이 non-linear submanifold로 확장될 수 있다고 추측한다.

4. Experiments

4.1 Experiments on Synthetic Data

The effect of $t_0$

FLIPD는 모든 이론들이 $\delta \rightarrow -\infty$ regime을 hold하기 때문에 $t_0$을 0에 가깝게 셋팅하는 것을 필요로 한다. Low-dimensional manifold에 fitting된 DMs가 $t_0 \searrow 0 $에 따라 수치적으로 unstable score $s(\cdot, t_0)$를 보인다는 사실은 FLIPD를 무효화하지는 않지만, 이는 $t_0$의 선택에 민감할 수 있다는 것을 제안한다. 그래서 본 논문에서는 $(0,1)$에서의 다양한 $t_0$에 대해서 $FLIPD(x,t_0)$에 대한 $t_0$의 효과를 실험한다.

1. $\mathbb{R}^10$에 embedding된 2,4,8 차원의 세 개의 isotropic Gaussians의 mixture인 분포
2. $\mathbb{R}^3$에 embedding된 2d torus와 1d circle의 mixture인 분포

위 두가지 분포에 대해서 DM을 학습시킨다.

$FLIPD(x,t_0)$는 $t_0 = 0$에서 부정확하지만, 금방 모든 datapoint에서 true LID 주변에서 안정화된다.

FLIPD is a multiscale estimator

Figure 2(b)에서 $t_0 = 0.65$에서 두 번째 knee가 발생한다. 이는 torus가 멀리서 봤을 때 1d circle처럼 보이기 때문이라고 한다. $t_0$의 값이 커지는 것은 manifold를 더 멀리서 보는 것과 대응되는 것이라고 한다.

Finding knees

Maximum curvature의 point를 찾는 knee detection algorithm인 kneedle을 활용한다.

Results

MAE에서 FLIPD가 model-based estimator에서 가장 좋은 경향을 보였고, 특히 dimension이 높아질 수록 그런 경향이 더 보였다. model-free baseline들이 간단한 시나리오에서 좋은 결과를 보였지만, LID가 증가하거나 data manifold에서 non-linearity가 더 도입될수록 unreliable results를 생성한다.

4.2 Experiments with Fully-Connected & UNet Architectures on Image Data

5. Conclusions, Limitations

비록 FLIPD가 synthetic benchmarks에 대해서 완벽한 LID estimates를 생성했지만, 이미지에 대한 network architecture의 선택에 관한 instability는 놀랍고 이는 이러한 선택에 대해 강하게 의존하는 LID estimates를 보여줬다. 본 논문에서 이러한 것들을 limitation이라고 언급하지만, 그럼에도 불구하고 FLIPD는 architecture와 관계없이 복잡성에 대한 의미있는 measure를 제공한다.