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' 논문을 기반으로 한다. 하나의 글로 작성하면 내용이 너무 길어질 것 같아서 '(1) Kernel Method', '(2) Neural Tangent Kernel', 그리고 '(3) Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains'의 세 부분으로 나눠서 작성하고자 한다. 이 글은 (1)을 설명한 지난 글에 이어 (2)에 관한 내용을 작성하고자 한다.

 

Neural Tangent Kernel에 관해 설명하기 전에 kernel의 정의와 kernel method에 관해 궁금하면 지난 글을 참고하면 된다.

 

(1) Kernel Method

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

 

 

 

NTK(Neural Tangent Kernel)

 

Neural Network와 Loss Function

 

 

지난 글에서 linear regression을 통해 예측하려는 예시에서는 함수가 모델의 파라미터 $w$에 관해 linear한 경우였다. 그러나 일반적인 neural network에서 추정하려고 하는 함수 $f$가 네트워크의 파라미터 $w$에 관해 non-linear 하면 어떻게 될까? Loss function을 지난 글의 예시처럼 MSE라고 가정하고 이에 관한 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) \end{align}$$

 

함수 $f$가 선형 모델일 때는 $\nabla_w f(w_t, x_i)$가 $x_i$로 깔끔히 정리되어 loss function의 gradient인 $\nabla_w \mathcal{L}(w_t)$가 $w_t$에 dependent 하지 않았지만, 여기서는 $f$를 non-linear 함수라고 가정했으므로 저 위의 식에서 더 이상 간단히 나타낼 수 없다. 그러므로 loss function의 gradient가 $w_t$에 dependent 하고, 결국 $t$에 따라 변하게 되면서 static 하지 않게 된다. 앞서 우리는 이를 static 하게 만들어야 loss function이 시간과 관계없이 어떻게 변하는지를 알 수 있다고 했는데, 여기서 어떻게 하면 loss function의 gradient를 static 하게 만들 수 있을까? 이를 가능하게 만들 수 있는 한 가지 전제를 설정해 보자.

 

 

 

Infinite-width Network에서의 NTK

 

Lazy Training

 

지난 글에서 설명했듯이 우리는 한 가지 중요한 성질을 알았다. 바로 너비가 무한히 큰(Infinite-width) 또는 데이터에 관해 충분히 너비가 넓은(sufficiently large) neural network에서는 $w_t$가 시간 $t$에 따라 거의 변하지 않는다는 것이며, 이를 논문에서는 'almost static' 하다고 표현한다. 즉, 모델의 파라미터가 처음 초기값인 $w_0$에서 시간에 따라 거의 변하지 않는다는 것인데, 이를 lazy training이라고 한다. 

다시 말해, 일반적인 네트워크 모델이라면 loss function의 gradient가 static 하지 않은데, 그 네트워크의 너비가 충분히 크면 $w$가 아주 천천히 변한다. 그러면 $w$가 그 초기 상태였던 $w_0$에서 거의 변하지 않으므로 "$w_0$에 관한 First Order Taylor Approximation을 써 보면 어떨까?"하는 아이디어에서 출발해 볼 수 있다.

 

 

[출처] https://rajatvd.github.io/NTK/

 

실제로 위의 그림은 어떤 task에 관해 학습을 수행하는 sufficiently large width를 지닌 모델의 파라미터가 시간에 따라 어떻게 변하는지를 행렬로 표현한 것이다. 각 칸의 색은 대응되는 weight의 크기를 뜻하는데, 시간에 따라서 매우 느리게 weight의 값이 변한다는 사실을 확인할 수 있다. 이러한 현상을 lazy training이라고 한다.

 

 

Infinite-width Network의 NTK

그래서 lazy training이라는 성질을 통해 $w$를 $w_0$ First Order Taylor Approximation을 적용하면 다음과 같은 식을 얻을 수 있다.

 

$$f(w, x) \approx f(w_0, x) + \nabla_w f(w_0, x)^{\intercal}(w - w_0)$$

 

위의 식을 자세히 들여다 보면 neural network이 학습하는 function인 $f$가 $w$에 관해서는 linear 하지만 $x$에 관해서는 linear 하지 않다. $\nabla_w f(w_0, x)$에서 $w_0$는 어떤 정해진 초기 값이므로 $w$에 관한 항이 되지 않는다. 그런데 이러한 역할을 하는 함수가 낯이 익다. 지난 글에서 kernel의 mapping function도 input에 관해서는 linear 하지 않지만, 이를 mapping function을 통해 고차원으로 보낸 결과에 관해서는 linear 했다. 이는 mapping function 자체가 그러한 역할을 하기 때문이다.

 

$$\phi(x) \triangleq \nabla_w f(w_0, x)$$

 

즉, 파라미터 $w$에 관한 $f$의 gradient 자체를 kernel의 mapping function으로 갖으므로 'neural tangent kernel'라는 이름을 붙인 것이다. 비문이어서 개인적으로도 이해하기 좀 어려웠는데, 쉽게 말해 함수 $f$를 추정하는 식을 근사해서 구해보니 그 파라미터의 초기 weight에 관한 $f$의 gradient를 kernel로 해석할 수 있으므로 neural tangent kernel이라고 부른다. 그러한 NTK가 neural network의 function $f$를 추정하는 데 있어서 중요하다고 정리할 수 있다.

 

 

$$K(x, x') = \left\langle \phi(x), \phi(x') \right\rangle = \left\langle \nabla_w f(w_0, x), \nabla_w f(w_0, x') \right\rangle$$

 

지난 글에서 확인한 kernel의 정의에 의해 NTK는 두 임의의 데이터 샘플에 관해 neural network가 학습하는 함수 $f$의 파라미터 $w$에 관한 gradient의 similarity를 계산하는 식으로 정리할 수 있다.

 

여기서 $\phi(x)$는 $w$의 초기 값인 $w_0$에 관해서 evaluate 하는 mapping function이므로, 근사한 식 $f(w, x)$는 시간 $t$에 invariant 하면서 동시에 convex optimization을 수행할 수 있다. 이를 통해 neural network의 학습 과정을 마치 kernel function을 찾는 것처럼 고려하여 linear하게 분석할 수 있으며, 이는 다시 말해 neural network가 non-linear 해도 마치 linear model인 것처럼 해석할 수 있다. 그래서 NTK만 구하면 무한히 큰 width를 지닌 neural network를 굳이 학습할 필요 없이 network의 성능에 관한 예측이 가능하다.

 

1-hidden Layer를 가진 Network의 NTK

 

직접 neural network를 학습하지 않고도 network 모델이 어떠한 방향으로 학습이 이루어지는지를 예측할 수 있다는 건데, 하나의 hidden layer를 지니는 간단한 neural network를 예시로 살펴보자. 위와 같이 input layer 1개, output layer 1개, 그리고 width가 $m$인($m$개의 neuron으로 구성된) hidden layer 1개로 구성된 network가 존재할 때, 파라미터와 activation 함수를 다음과 같이 표기한다고 가정한다.

 

$w$: $a_i$와 $b_i$의 집합

$\sigma$: activation function인 ReLU

 

$m$이 커져서 infinite 해지면 neural tangent kernel을 사용함으로써 학습 과정이 linear model처럼 근사될 수 있다는 것인데, 이는 다시 말해 학습 과정에서 파라미터들의 변화가 linear해짐을 의미한다. 정말 그런지 확인하기 전에 앞서 NTK를 알면 그 모델이 어떠한 방향으로 학습이 진행되는지를 예측할 수 있다고 했는데, 먼저 이에 관해서 알아보자.

 

 

 

$$\hat{y} = f_m(w, x) = \frac{1}{\sqrt{m}} \sum_{i=1}^m b_i \sigma(a_i^{\intercal}x)$$

 

모델의 함수 $f_m$(width가 $m$인 network에 관한 함수)에 관한 데이터 $x$의 예측값을 정리한 식이다. 여기서 $\frac{1}{\sqrt{m}$는 뒤에서 기댓값으로 식을 정리할 때 편하게 하려고 곱하는 항이다.

 

 

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

 

Loss function은 위와 같이 MSE로 정의한다고 가정한다.

 

 

$$\nabla \mathcal{L}(w)= \sum_{i=1}^n (y_i - f_m(w, x_i)) \nabla_w f_m(w, x_i)$$

 

이 loss function에 관하여 파라미터 $w$에 대한 gradient를 구한다. 그런데 위의 network에서 파라미터 $w$는 앞단에 있는 $a_i$와 뒤에 있는 $b_i$ 파라미터들의 집합이므로 각각의 파라미터 종류에 따라 $f_m$의 gradient를 따로 정리해 본다.

 

 

$$\nabla_{a_i} f_m(w, x) = \frac{1}{\sqrt{m}} b_i \sigma^{\prime}(a_i^{\intercal}x) (a_i^{\intercal}x)^{\prime} =\frac{1}{\sqrt{m}} b_i \sigma^{\prime}(a_i^{\intercal}x) x$$

$$\nabla_{b_i} f_m(w, x) = \frac{1}{\sqrt{m}} \sigma(a_i^{\intercal}x)$$

 

각각 파라미터 $a_i$와 $b_i$에 대한 $f_m$의 gradient를 정리한 결과다.

 

 

 

그런데 이 network의 width가 sufficiently large 하다고 가정했으므로 lazy training 성질에 의해 $w$를 $w_0$와 같게끔 생각하여 다음과 같이 식을 정리할 수 있다.

 

$$\nabla_{a_i} f_m(w, x) = \nabla_{a_i} f_m(w_0, x) = \frac{1}{\sqrt{m}} b_i \sigma^{\prime}(a_i^{\intercal}x) (a_i^{\intercal}x)^{\prime} =\frac{1}{\sqrt{m}} b_i \sigma^{\prime}(a_i^{\intercal}x) x$$

$$\nabla_{b_i} f_m(w, x) = \nabla_{b_i} f_m(w_0, x) = \frac{1}{\sqrt{m}} \sigma(a_i^{\intercal}x)$$

 

개인적으로 여기서 의문이 드는 점은 저기서 $w$와 $w_0$를 같게끔 했으니 $w$에 속하는 파라미터 $a_i$와 $b_i$도 각각의 파라미터에 관한 초기값으로 고려해도 되는 건지였다. 그렇게 이해하는 게 맞다고 생각했는데, 이 부분은 더 조사를 한 후 추후 보강하도록 하겠다.

 

 

 

이제 파라미터 $a_i$와 $b_i$에 대한 kernel을 다음과 같이 정리할 수 있다. 

 

$$K_m^{(a)} (x, x') = \left\langle \nabla_{a_i} f_m (w_0, x) , \nabla_{a_i} f_m (w_0, x') \right\rangle$$

$$K_m^{(b)} (x, x') = \left\langle \nabla_{b_i} f_m (w_0, x) , \nabla_{b_i} f_m (w_0, x') \right\rangle$$

 

그런데 $w$는 $a_i$와 $b_i$의 집합이므로 concatenation을 사용하여 다음과 같이 쓸 수 있다.

 

$$K_m(x, x') = \left\langle \nabla_{a_i} f_m (w_0, x) \oplus \nabla_{b_i} f_m (w_0, x), \nabla_{a_i} f_m (w_0, x') \oplus \nabla_{b_i} f_m (w_0, x') \right\rangle$$

 

Kernel 식을 계산하는 건 inner product이므로 $a_i$와 $b_i$에 대한 $f_m$의 gradient를 내적하는 건 결과적으로 다음 식을 계산하는 것과 같다.

 

$$K_m(x, x') = K_m^{(a)} (x, x') + K_m^{(b)} (x, x')$$

 

 

 

위에서 정리한 $a_i$와 $b_i$에 대한 kernel을 풀어서 정리해보면 다음과 같다.

 

$$\begin{align}K_m^{(a)} (x, x') &= \left\langle \nabla_{a_i} f_m (w_0, x) , \nabla_{a_i} f_m (w_0, x') \right\rangle\\
&= \frac{1}{m} \sum_{i=1}^m b_i^2 \sigma^{\prime} (a_i^{\intercal}x) \sigma^{\prime}(a_i^{\intercal}x') \left\langle x, x' \right\rangle\\&= \mathbb{E}\left[ b_i^2  \sigma^{\prime} (a_i^{\intercal}x) \sigma^{\prime}(a_i^{\intercal}x')\left\langle x, x' \right\rangle \right]\end{align}$$

$$\begin{align}K_m^{(b)} (x, x') &= \left\langle \nabla_{b_i} f_m (w_0, x) , \nabla_{b_i} f_m (w_0, x') \right\rangle\\&= \frac{1}{m} \sum_{i=1}^m  \sigma (a_i^{\intercal}x) \sigma(a_i^{\intercal}x')\\&= \mathbb{E}\left[ \sigma (a_i^{\intercal}x) \sigma(a_i^{\intercal}x') \right]\end{align}$$

 

이제까지는 network의 width가 finite한 $m$일 때를 가정했는데, 만약에 $m$이 무한대에 가까워지면 어떻게 될까? Neuron의 개수가 매우 많으므로 직접 일일이 계산할 수는 없지만, 앞서 구한 kernel 식을 바탕으로 어떠한 결과로 수렴할지는 예측할 수 있다.

 

초기 값으로 고려해야 하는지는 앞서 말한 바처럼 확실치는 않으나, 만약에 $a_i$와 $b_i$에 관한 초기 값을 분포로 고려하면 $m$을 무한대로 보냈을 때 각각의 파라미터 종류에 대한 kernel을 아래와 같이 정리할 수 있다.

 

$$K_m^{(a)} (x, x') \xrightarrow{m \rightarrow \inf} K^{(a)} (x, x')$$

$$K_m^{(b)} (x, x') \xrightarrow{m \rightarrow \inf} K^{(b)} (x, x')$$

 

큰 수의 법칙에 의해 $m$이 무한대에 가까워질수록 각각의 파라미터 종류에 관해 앞서 정리했던 식의 기댓값으로 수렴한다는 것을 예상할 수 있다. 즉, 직접 일일이 학습 과정을 거치거나 연산을 수행하지 않아도 무한대의 width를 가진 neural network의 kernel을 구했다고 볼 수 있으며, 모델의 학습 과정이 어떻게 동작하는지를 이론적으로 보였다고 요약할 수 있다.

 

 

 

NTK와 Neural Network의 학습

 

데이터에 관해 sufficiently large한 network에 관해서 NTK를 통해 파라미터 weight을 optimization 하는 과정을 분석할 수 있고, 그러한 NTK가 모델의 weight을 학습하는 과정을 linear하게 해석할 수 있도록 한다는 것까지는 알아봤다. 그런데 이렇게 학습할 때 과연 $w$이 정말 실제 함수를 찾을 수 있도록 수렴하는 걸 장담할 수 있는지는 보이지를 못했다. 즉, infinite width를 지닌 neural network에 관해 linear function의 존재성(existence)은 보였지만, 이를 학습해서 과연 그러한 function을 찾을 수 있는지 그 수렴 가능성을 보이는 것도 중요하다. 이를 확인하기 위해 모델의 파라미터가 학습 과정에서 어떻게 변하는지를 알아볼 필요가 있는데, gradient flow dynamics를 사용하요 알아보자. 막 이것저것 설명하다보니 흐름을 놓칠 수 있는데, 앞서 확인한 건 모델의 파라미터를 최적화 하는 과정을 NTK를 통해 해석할 수 있다는 얘기이고, 여기부터는 gradient flow dynamics를 사용하여 실제로 파라미터가 학습 과정에서 어떻게 바뀌는지를 알아보고자 하는 것이다.

 

 

Gradient Flow Dynamics

 

$$w(t+1) = w(t) - \eta_t \nabla_w \mathcal{L}(w(t))$$

 

파라미터 weight이 시간이 따라서 어떻게 변하는지를 gradient descent로 정리한 식이다. $w(t)$를 좌변으로 보내고 양변을 learning rate $\eta_t$로 나눠주면 다음과 같다.

 

$$w(t+1) - w(t) = - \eta_t \nabla_w \mathcal{L}(w(t))$$

$$\frac{w(t+1) - w(t)}{\eta_t} = - \nabla_w \mathcal{L}(w(t))$$

 

그런데 learning rate $eta_t$를 0에 가까운 매우 작은 값으로 고려하면 좌변을 weight $w$가 시간에 따라 변하는 어떠한 gradient 함수로 볼 수 있다. 마치 $w$에 대한 derivative term으로 고려하는 것인데, 이를 다음과 같이 정리할 수 있다.

 

$$\frac{\operatorname{d}\!w(t)}{\operatorname{d}\!t} = - \nabla_w \mathcal{L} (w(t))$$

 

이러한 분석 방법을 asymptotic regime라고 하는데, 이러한 방법이 loss function의 특성을 지니면서도 이를 단순화한 형태에서 분석 가능하다는 것이다.

 

 

 

$$\mathcal{L}(w) = \frac{1}{2} \| \hat{y}(w) - y \|^2$$

$$\nabla_w \mathcal{L}(w(t)) = \nabla_w \hat{y}(w(t)) \left(\hat{y}(w(t)) - y \right)$$

 

Loss function이 MSE라고 가정하고 $w$에 관한 gradient를 구하면 다음과 같이 정리할 수 있다.

 

 

$$\frac{\operatorname{d}\!w(t)}{\operatorname{d}\!t} = - \nabla_w \hat{y}(w(t)) \left(\hat{y}(w(t)) - y \right)$$

 

그런데 앞서 $\frac{\operatorname{d}\!w(t)}{\operatorname{d}\!t} = - \nabla_w \mathcal{L} (w(t))$라고 정리했으므로 우변을 풀어서 쓴 식으로 바꿔주면 위와 같다. 이것이 weight $w$가 시간 $t$에 따라 $w$-space에서 어떻게 변하는지를 보여준다.

그러나 여기까지는 별다른 소득이 없어 보이는데, weight $w$ 자체를 보기 보다는 다시 앞서 봤던 것처럼 network가 학습하려는 함수 $f$가 시간에 따라 어떻게 변하는지인 $\frac{\operatorname{d}\!f(w(t))}{\operatorname{d}\!t}$를 확인해 보자.

 

 

$$\hat{y}(w(t)) = f(w(t))$$

$$\frac{\operatorname{d}\!f(w(t))}{\operatorname{d}\!t} = \frac{\operatorname{d}\!\hat{y}(w(t))}{\operatorname{d}\!t}$$

 

데이터에 관해 모델이 예측한 값이 결국 그 함숫값이므로 이를 다시 정리하면 위처럼 된다. 이제 이를 풀어서 쓰면 아래와 같은 과정으로 정리할 수 있다.

 

$$\begin{align}\frac{\operatorname{d}\!\hat{y}(w(t))}{\operatorname{d}\!t} &= \nabla_w \hat{y}(w(t))^{\intercal}\frac{\operatorname{d}\!w(t)}{\operatorname{d}\!t}\\&= -\nabla_w \hat{y}(w(t))^{\intercal}\nabla_w\hat{y}(w(t)) \left(\hat{y}(w(t)) - y \right)\\&= -\nabla_w f(w)^{\intercal} f(w) \left(\hat{y}(w(t)) - y \right)\end{align}

 

정리해서 봤더니 $\nabla_w f(w)^{\intercal} f(w)$라는 익숙한 표현이 나온다. 이는 gradient끼리 내적한 것인데, 앞서 NTK에서도 확인했다시피 gradient끼리 내적한 걸 kernel로 볼 수 있다. 이를 kernel로 고려하여 식을 정리하면 아래와 같다.

 

$$\frac{\operatorname{d}\!\hat{y}(w(t))}{\operatorname{d}\!t} = - K(w) (\hat{y}(w(t)) - y)$$

 

위의 식이 결국 $\hat{y}$가 gradient descent의 trajectory에 따라 어떻게 변하는지를 보여준다.

 

 

 

그런데 neural network의 width가 충분히 크면 lazy training에 의해 $w \approx w_0$라고 했으므로 $K(w) \approx K (w_0)$이고, 이를 바탕으로 위의 식을 다시 정리하면 아래와 같다.

 

$$\frac{\operatorname{d}\!\hat{y}(w(t))}{\operatorname{d}\!t} = - K(w_0) (\hat{y}(w(t)) - y)$$

 

여기까지도 우리가 목표로 했던 $w$의 최적화를 통한 함수의 수렴 가능성을 보이진 못한다. 그래서 ground truth의 $\hat{y}$와 예측 값 $y$간의 error인 $\hat{y} - y$가 시간에 따라서 어떻게 변하는지를 알아보자.

 

 

Infinite-width Network의 Convergence

 

$$u = \hat{y} - y$$

$$\frac{\operatorname{d}\!u}{\operatorname{d}\!t} \approx - K(w_0) \cdot u$$

 

Ground truth의 $\hat{y}$와 예측 값 $y$간의 error인 $\hat{y} - y$를 $u$로 정의해 보자. 그러면 $u$에 관해 시간 $t$에 대한 derivative는 위와 같이 정리할 수 있다. 그런데 위의 식은 1계 ODE 문제이므로 공학수학에서 배웠던 변수분리법에 의해 양변을 각각의 변수로 정리한다.

 

$$\begin{align}\int \frac{1}{u} \operatorname{d}\!u &= \int -K(w_0) \operatorname{d}\!t\\\ln (u) &= -K(w_0)t + C\\u(t) &= C e^{-K(w_0)t}\\\end{align}$$

 

1계 ODE을 풀었는데 상수 $C$가 남는다. 이 상수를 $u$에 관한 항으로 바꾸면 다음과 같이 정리할 수 있다.

 

$$u(0) = C \cdot 1$$

$$u(t) = u(0) \, e^{-K(w_0)t}$$

 

여기서 모델의 파라미터가 실제 데이터에 근사하게끔 수렴하는지를 알 수 있다. 즉, 학습을 통해 데이터에 잘 근사할 수 있는 함수 $f$를 찾을 수 있다는 뜻이다. 이 부분이 이해가 잘 안 간다면 위의 $e^{-K(w_0)t}$에 관한 지수 함수 그래프를 참고해 보면 된다. $t$가 무한대에 가까워질수록 $u$는 0으로 수렴하고, 이는 결국 $\hat{y}$가 $y$로 수렴한다는 걸 의미한다.

 

여기까지 보았던 내용을 정리하면, infinite width를 지닌 network 모델의 파라미터를 학습하는 데 있어서 NTK라는 것을 확인했고 이 NTK가 파라미터를 학습하는 과정을 linear 하게 분석하게 만든다는 것이었다. 또한 그러한 모델의 파라미터는 학습을 수행하면 결국 실제 데이터에 수렴하게 된다는 사실도 확인했다. 개인적으로 느끼기에는 수학적으로 정리할 내용이 적지 않고 이러한 사고 흐름을 따라가는 것도 쉽지가 않았다. 괜히 논문 출판 당시 리뷰에 '수학적으로 이해하기 난해하다'라는 평이 있었던 게 아닌 거 같다.

 

다음 글에서는 마지막으로 NTK의 positive definite한 특성을 통해 왜 모델이 학습할 때 high frequency에 관해서는 집중을 잘 하지 않는 것인지를 확인해보고, 이를 극복하기 위해 제안한 input encoding 방법인 Fourier feature를 사용한 positional encoding에 관해 정리해 보고자 한다.

 

 

[이전 글]

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

 

[다음 글]

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

 

 

출처
이 내용을 공부하고 글로 정리하는 데 있어서 큰 도움이 된 영상과 자료들이다.

1. https://youtu.be/9Jx2ulS1sAk, PR-374: Fourier Features Let Networks Learn High-Frequency Functions in Low Dimensional Domains
2. http://jnwoo.com/archive/6, Neural Tangent Kernel
3. https://rajatvd.github.io/NTK, Understanding the Neural Tangent Kernel