AI/CV

Neural Tangent Kernel과 Fourier Features를 사용한 Positional Encoding (1) - Kernel Method

Positional encoding은 AI 모델링 분야에서 많이 쓰이는 기법 중 하나이다. 대표적으로 transformer에서도 데이터를 병렬 처리하여 학습하는 단점을 보완하여 데이터를 구성하는 각 token에 위치 정보를 부여하고자 input을 sinusoidal 등 어떠한 함수에 통과시켜 모델에 넣는 과정을 positional encoding이라고 한다. 컴퓨터 비전에서는 목적은 다르지만 이와 유사하게 사용되는 Fourier feature를 이용한 positional encoding이 존재하는데, 좀 더 high frequency 정보를 잘 학습할 수 있도록 하기 위함이다. 그러나 개인적으로 여태까지 이를 여과없이 단지 "적용하면 좋다"라는 '카더라' 식의 얘기로만 이해했지, 구체적으로 왜 그런 건지 그 근거는 모르고 있었다. 막연하게 '그렇더라'고 여기기 보다는 좀 더 깊게 파 볼 필요가 있어서 현재 근무 중인 연구실의 세미나 발표 주제로 준비하기로 했다. 하지만 수학적 논리와 지식이 부족한 나에게는 그 기저에 있는 내용들이 많이 어렵게 느껴졌고, 아직도 완전히 이해했다고 절대 말할 수 없다. 다행히 여러 사람들이 좋은 자료를 남겨줘서 이를 참고하면서 학부생 수준으로서 조금씩 이해해 보려고 노력했고, 이 글도 그러한 관점에서 논문의 흐름대로 수학적으로 완벽히 증명하기 보다는 그러한 개념이 등장하게 된 배경과 이를 증명하는 과정에서의 사고 흐름에 집중하고자 한다. 엄밀한 증명보다는 직관적인 이해를 목적으로 정리하는 글이므로 흐름을 좇아가는 데 주목해 보자.
이 글은 'Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains' 논문과 이를 이해하는 데 필요한 선수 내용을 소개한 'Neural Tangent Kernel: Convergence and Generalization in Neural Networks' 논문을 기반으로 한다. 글의 내용이 너무 길어질 것 같아서 '(1) Kernel Method', '(2) Neural Tangent Kernel', 그리고 '(3) Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains'의 세 부분으로 나눠서 작성하고자 한다.

Introduction

[출처] Ben Midenhall, NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis / Katja Schwarz, GRAF: Generative Radiance Fields for 3D-Aware Image Synthesis

컴퓨터 비전에서 NeRF를 필두로 하여 Mip-NeRF, GRAF 등 3D-aware synthesis를 수행할 수 있는 neural network과 결합된 다양한 모델이 등장했는데, 적지 않은 논문들에서 모델의 input에 관해 Fourier feature로서의 positional encoding을 수행하는 것을 볼 수 있다. 이는 본 논문에서 input을 neural network에 넣기 전에 Fourier feature를 적용하는 것이 high frequency의 데이터 정보를 더 잘 잡아낼 수 있다는 사실을 보인 이유이다. 그런데 과연 왜 positional encoding을 하는 게 높은 주파수에 대응되는 데이터를 잘 학습할 수 있는 걸까? 이를 이해하려면 먼저 kernel method를 알아야 하고, neural network와 kernel이 어떠한 관계가 있어서 NTK(Neural Tangent Kernel)이라는 개념이 나온 것인지, NTK를 통해 왜 positional encoding 없이는 high frequency 데이터를 잘 학습하지 못하는 건지의 순서로 알아볼 필요가 있다.

그래서 앞서 preliminary로서 linear regression, mapping function, kernel method, 그리고 NTK에 관해 소개하고, 그 이후에 Fourier Feature에 관한 내용을 자세히 다루고자 한다.

Linear Regression

Linear Regression이란?

Linear regression은 실제 데이터에 잘 근사할 수 있는 함수 $f$를 linear function으로 가정하여 그 함수의 가중치를 찾는 과정으로 볼 수 있다. 그러므로 함수 $f$는 가중치 $w$에 관한 선형결합인 $f(w, x) = w^{\intercal}x$로 이루어지고, 적합한 가중치를 찾았을 때 임의의 데이터 $x_i$에 관해 대응되는 함수 값 $f(w, x_i)$가 바로 그 데이터에 관한 예측 값인 $\hat{y}_i$가 된다. 그러면 함수에서의 가중치인 $w$를 어떻게 최적화하는지가 중요한데, 일반적으로 이는 GD(Gradient Descent) 방법을 사용하여 해결한다.

GD를 수행하려면 실제 값과 예측 값 간의 차이를 정의하는 loss function이 존재해야 하는데, 간단하게 loss function을 MSE(Mean squared error)로 정의할 수 있고, 수식으로 정의하면 다음과 같다.

$$\mathcal{L}(w) = \frac{1}{2} \sum_{i=1}^n (y_i - \hat{y}_i)^2= \frac{1}{2} \sum_{i=1}^n (y_i - f(w,x_i))^2$$

이러한 loss function을 최소화할 수 있는 파라미터 $w$를 찾아야 하므로 이를 수학적으로 정리하면 $\min_w \mathcal{L}(w)$로 쓸 수 있다. 즉, 파라미터를 매번 epoch 또는 시간 $t$마다 업데이트 할 때 현재 그 파라미터에서 loss function을 최소화할 수 있는 방향인 gradient로 이동시켜야 하며, 이를 식으로 나타내면 $w(t + 1) = w(t) - \eta_t \nabla_w \mathcal{L}(w_t)$로 작성할 수 있다. 그런데 앞서 f(w, x) = w^{\intercal}x라고 가정했으므로 이를 파라미터 $w$에 관해 미분하면 $\nabla_w f(w_t, x_i) = \nabla_w (w^{\intercal}x_i) = x_i$로 쓸 수 있다.

다시 앞으로 와서 파라미터를 업데이트 하는 식인 $w(t + 1) = w(t) - \eta_t \nabla_w \mathcal{L}(w_t)$를 살펴 보자. 이를 계산하려면 loss function $\nabla_w \mathcal{L}(w_t)$에 관한 gradient를 구해야 하므로 다음과 같이 작성할 수 있다.

$$\begin{align} \nabla_w \mathcal{L}(w_t) &= \sum_{i=1}^n (y_i - f(w_t, x_i))\nabla_w f(w_t, x_i)\\ &= \sum_{i=1}^n (y_i - f(w_t, x_i)) x_i \end{align}$$

여기서 확인할 수 있는 사실은 xi에 관하여 $w_t$에 대한 함수 $f$의 gradient인 $\nabla_{w_t} f(w_t, x_i)$가 $w_t$에 dependent 하지 않다는 것이다. 즉, $x_i$에만 의존하고 시간 $t$에 dependent 하지 않아서 함수 $f$의 변화가 static 해지므로 파라미터를 optimization 하는 데 수월하다는 것이다. (물론 loss function의 $w$에 대한 gradient는 static 하지 않을 수 있다.)

Linear Regression의 한계

앞서 linear function $f$를 찾는 방법을 살펴 봤는데, 그러면 이러한 linear regression의 한계는 무엇일까? 바로 함수 $f$가 linear power만 지니고 있어서 복잡한 task를 수행하기 어렵다는 것이다.

예를 들어 보자. 위의 왼쪽 사진에서 노란색과 파란색으로 칠해진 두 group의 데이터에 관해 classification을 수행한다고 가정한다. 두 group을 구분지을 수 있는 boundary를 찾아야 하는데, 평면 상에서는 그러한 경계가 곡선의 형태를 띠고 있어서 linear function으로는 적절히 분류할 수가 없으며, 이를 'linearly inseparable'이라고 한다. 이를 해결하기 위한 한 가지 방법으로는 위의 오른쪽 사진에서처럼 데이터들을 고차원으로 보낸 후, 그 고차원 feature space 상에서 classifcation의 boundary를 결정지을 linear hyperplane을 찾는 것이다. 즉, 데이터들을 고차원으로 끌어 올림(lifting)으로써 저차원에서 linearly inseparable 했던 상태를 separable 하게 바꿀 수 있다는 것이다.

Mapping Function

Mapping Function이란?

데이터를 고차원으로 mapping 시켜서 non-linearity를 부여할 수 있도록 해야 하는데, 그러한 함수를 mapping function이라고 부르고, 여기서는 $\phi$라고 적겠다. Non-linearity를 부여하는 함수 $\phi$를 통해 저차원의 데이터를 linearly separable 할 수 있도록 더 고차원의 feature space로 보내는 것이다. Mapping function을 데이터에 적용한 후 함수 $f$를 작성하면 다음과 같다.

$$\begin{align} y = f(w, x) &= w^{\intercal}\phi(x)\\ &= w_1 \phi(x_1) + \cdots + w_m \phi(x_m) \end{align}$$

Linear regression과는 달리 찾으려고 하는 함수 $f$가 $x$에 관해서는 linear 하지 않지만, $w$에 관해서는 linear 하다는 걸 알 수 있다.

Mapping Function의 한계

Mapping function을 사용한 함수 $f$의 gradient descent를 구하는 것도 어렵지 않다. Loss function $\mathcal{L(w)}$와 이를 최소화할 수 있는 파라미터 $w$를 찾는 방법을 식으로 정리하면 다음과 같다.

$$\mathcal{L}(w) = \frac{1}{2} \sum_{i=1}^n (y_i - w^{\intercal}\phi(x_i))^2$$

$$\min_w\frac{1}{2} \sum_{i=1}^n (y_i - w^{\intercal}\phi(x_i))^2$$

그러나 단순히 mapping function을 사용하는 건 바람직하지 않을 수 있는데, 크게 두 가지의 한계가 존재해서다.

첫 번째는 feature map $\phi$를 어떻게 찾아야 하는지가 중요하다. 자신이 지닌 데이터에 적합한 $\phi$를 찾는 게 관건인데, 이를 위해 어떠한 non-linear function을 사용할지 사용자가 직접 골라야 한다. 이는 데이터에 잘 fit 되지 않는 함수를 찾을 수 있는 위험성을 지니고 있는 동시에 그러한 함수를 찾아야 한다는 번거로움이 존재한다.

두 번째는 데이터 $x$를 mapping function $\phi$를 적용하여 $\phi(x)$로 만든 후 어떠한 연산을 수행하게 될 텐데, 이때 끌어 올리려는 고차원의 차원 수 $D$가 너무 크면 메모리도 많이 들고 연산도 많이 필요하다. 모든 관측치에 관해 각각 고차원으로 mapping 하고 이에 관해 어떠한 연산(주로 내적)을 수행해야 하므로 computational resource가 늘어날 수 있다.

만약에 $x \in \mathbb{R}^d$인 데이터에 관해 mapping function $phi$를 적용하여 polynomial with $k$를 고차원에서 수행하면 complexity는 $O(D^k)$일 것이다. 그러나 후술할 polynomial kernel을 적용하면 $O(d^k)$ 안에 끝낼 수 있다. 이처럼 'kernel'이라는 것을 사용하면 전술한 mapping function의 단점을 극복할 수 있는데, 본격적으로 kernel을 알아보기 전에 polynomial kernel이 무엇인지 알아보면서 kernel의 특징을 직관적으로 이해해 보자.

Kernel

Polynomial Kernel

Polynomial kernel은 $(a \times b + r)^k$로 정의되는 어떠한 함수인데, 이 함수의 output은 두 관측 데이터 간의 relationship을 계산한다. 그런데 여기서 중요한 특징은 바로 직접 고차원으로 데이터를 변환하지 않고도 마치 데이터가 고차원에 있는 것처럼 고려하여 두 데이터 간의 relationship을 구할 수 있다는 것이다. 이를 시각적으로 잘 설명한 영상 link를 아래 첨부했다.

위의 영상에서 나오는 polynomial의 과정과 그 의미를 한 장의 그림으로 요약하면 위의 사진과 같다. $r=\frac{1}{2}$이고 $k=2$일 때의 $(a \times b + r)^k$를 계산하는 과정을 풀어쓰면 다음과 같다.

$$\begin{align}(a \times b + \frac{1}{2})^2 &= (a \times b + \frac{1}{2})(a \times b + \frac{1}{2})\\&= a^2 b^2 + ab + (\frac{1}{2})^2 = ab + a^2b^2 + \frac{1}{4}\\&= \left( a, a^2, \frac{1}{2} \right) \cdot \left( b, b^2, \frac{1}{2} \right) \end{align}$$

즉, $(a \times b + r)^k$ 식을 단순히 계산하는 것만으로도 1차원의 데이터를 2차원으로 보내서 그 2차원 공간에서의 관계를 구할 수 있다는 의미이다. 위의 식을 살펴 보면 1차원의 데이터를 2차원으로 각각의 데이터를 보냈을 때의 2차원에서의 좌표와 그 좌표들의 dot product로 정리된다는 사실을 알 수 있다. 핵심은 직접 고차원의 좌표로 바꿔서 내적 하지 않아도 polynomial kernel만으로 계산이 가능하다는 것이다.

Kernel Function and Kernel Method

앞서 polynomial kernel의 특징을 통해 kernel이 어떠한 역할을 하는지 조금씩 감이 잡혔을 것이다. Kernel function은 mapping function으로 데이터를 고차원으로 보내지 않고도 고차원 공간에서의 데이터 간의 관계를 output으로 내 보내는 함수이며, kernel function을 사용하는 방법을 kernel method 또는 kernel trick이라고 한다. 좀 더 수학적으로 kernel을 정의하면 다음과 같다.

$$\begin{align}\phi: \quad &\mathbb{R}^{1\times n_0} \rightarrow \mathbb{R}^{n_L\times 1}\\ K: \quad &\mathbb{R}^{n_0} \times \mathbb{R}^{n_0} \rightarrow \mathbb{R}^{n_L \times n_L}\\ &(x, x') \longmapsto \left\langle \phi(x), \phi(x') \right\rangle \end{align}$$

위의 수학적인 표현을 해석하면 $1 \times n_0$ 차원의 두 벡터 $x$, $x'$를 tuple 형태의 input으로 받아서 $x$, $x'$ 각각의 벡터에 관해 $n_L \times 1$ 차원의 feature vector를 구한 후, 이들의 cross similarity를 표현한 $n_L \times n_L$ 행렬을 output으로 내는 함수라고 할 수 있다.

Kernel Method 사용 이유

그러면 이러한 kernel method가 mapping function에 비해 어떠한 장점을 지니고 있는지를 정리해 보자.

Kernel function은 수학적으로 데이터의 고차원 공간에서 두 데이터 간의 내적으로 해석할 수 있는데, 대부분의 ML algorithm에서는 내적만 알아도 되는 경우가 많다고 한다. 사실 이 부분이 크게 와닿지는 않았는데, 실제로 AI 모델링을 하면서 모델이 어떠한 feature를 학습하는 데 있어서 concatenation도 있지만 inner product로 처리하는 경우가 많다. 그러한 맥락에서 나오는 성질인지는 확실치는 모르겠지만, 여하튼 자명하다고 생각하고 넘어가면서 kernel function의 내적 표현이 결과적으로 ML algorithm에 잘 적용될 수 있다고 받아들였다. (나중에 기회가 되면 그 근거를 찾아 좀 더 확실하게 작성해 보겠다.) Kernel method의 장점이라고 하면 mapping function $\phi$가 무엇인지 굳이 explicit 하게 알지 않아도 feature space에서의 관계를 구할 수 있다는 것이다. 또한 입력으로 온 각각의 벡터를 고차원으로 올리지 않아도 그 입력 벡터의 차원에서 연산이 가능하다는 점이다. 즉, 데이터에 관해 $\phi$를 따로 통과시킬 필요가 없다.

그런데 여기서 의문이 드는 게, 도대체 neural network와 kernel이 어떠한 관계에 있기에 Neural Tangent Kernel이라는 말이 나왔냐는 것이다. 이에 관해서는 한 가지 짚고 넘어가야 할 점이 있는데, 바로 어떤 한 neural network에 관해 estimated function을 학습하는 과정은 결과적으로 kernel regression과 같다는 것이다. 단, 전제는 그 neural network가 infinite-width 또는 데이터에 관해 sufficiently large width를 지니고 있어야 한다는 것이다. 이러한 사실을 어떻게 받아들일 수 있는지는 다음 글에서 Neural Tangent Kernel을 설명하면서 이어가고자 한다.

[다음 글]

Neural Tangent Kernel과 Fourier Features를 사용한 Positional Encoding (2) - Neural Tangent Kernel

이 글 시리즈는 'Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains' 논문과 이를 이해하는 데 필요한 선수 내용을 소개한 'Neural Tangent Kernel: Convergence and Generalization in Neural Networks'

glanceyes.com

Neural Tangent Kernel과 Fourier Features를 사용한 Positional Encoding (3) - Fourier Features