Neural Tangent Kernel 리뷰

이번에 리뷰할  Neural Tangent Kernel (NTK) 논문은 NIPS 2018에 실린 논문으로 많은 인용수를 자랑하는 파급력 높은 논문입니다. 하지만 논문을 이해하려면 수학적 배경지식이 많이 필요해서 읽기가 어렵습니다. 다행히 설명을 잘해놓은 외국 블로그([2])가 있어 많은 참고를 하였습니다.

 

이 논문은 2 hidden layer와 infinite nodes의 neural network는 linear model로 근사하여 생각할 수 있고, 그 덕분에 문제를 convex 하게 만들어 해가 반드시 존재한다는 것을 보여줍니다. 어떻게 linear model로 근사하여 생각할 수 있다는 것인지 차근차근 알아보겠습니다.

 

Taylor Expansion

Taylor expansion은 매우 작은 영역에서 함수를 근사하는 방법입니다. Update된 parameter $ w $와 initial parameter $ w_{0} $에 대해 1차 derivative까지 근사하여 taylor expansion 하면 아래와 같이 나타낼 수 있습니다.

$$ y(w) \approx y(w_{0}) + \triangledown_{w} y(w_{0})^{T} (w-w_{0}) $$

위에서 언급했듯 taylor expansion은 매우 작은 영역에서 잘 성립하는 방법입니다. 위의 식이 성립하려면 initial parameter $w_{0}$와 $w$ 사이의 차이가 별로 나지 않아야 합니다. 흥미롭게도 $w$와 $w_{0}$는 network width가 커질수록(node가 많아질수록) 그 차이가 줄어듭니다. 이는 직관적으로 당연한 결과인데, parameter가 많아질수록 조금의 변화로 큰 변화를 일으킬 수 있기 때문입니다. 이에 대한 실험 결과가 블로그([2])에 나와있어 figure 1에 가져왔습니다.

Figure 1. Width가 커질수록 loss가 빠르게 줄어들고(left) $ w $와 $ w_{0} $의 차이가 작은 것을(right) 볼 수 있다.

왼쪽 그림의 x축은 업데이트 횟수이고 y축은 training loss이며 width가 다를 때 어떻게 변화하는지 보여줍니다. Width가 클수록 적은 step만으로 training loss가 빠르게 작아짐을 볼 수 있습니다. 오른쪽 그림의 y축은 $ w_{0} $로부터 n번 update 한 parameter $ w_{n} $까지의 거리를 나타냅니다. Width가 클수록 움직인 거리가 작아짐을 볼 수 있습니다.

 

Gradient Flow

Width가 infinite이면 taylor expansion으로 근사하는 것을 사용할 수 있다는 것을 알 수 있습니다. 이제 gradient flow를 전개하여 width가 infinite이면 해가 반드시 존재한다는 것을 보이겠습니다.

Parameter $w$가 업데이트 되는 것은 다음과 같이 표현할 수 있습니다. ($\eta$는 learning rate)

$$ w_{k+1} = w_{k} - \eta \triangledown_{w} L(w_{k}) $$

$$ \frac{w_{k+1}-w_{k}}{\eta} = -\triangledown_{w} L(w_{k})  $$

$ \eta $(learning rate)가 매우 작아 0에 근사하여 continuous하게 update 한다고 생각하면 다음과 같이 표현할 수 있습니다.

$$ \frac {dw(t)}{dt} = -\triangledown_{w} L(w(t)) $$ 

여기서 dt는 보통 time에 대한 변화량을 나타내는데, continuous 하게 변화한다고 생각하여 그렇게 표현한 것뿐입니다. 실제로는 update 횟수라고 생각하면 됩니다. 이제 $\frac {dw(t)}{dt} = \dot {w} $ 로 나타내서 표현하겠습니다.

$$ \dot{w} = -\triangledown_{w} L(w(t)) $$

Loss function은 mean squared error로 가정하여 gradient를 계산하면 다음과 같습니다. ($\bar {y}$는 label)

$$ \dot{w} = -\triangledown y(w)(y(w)-\bar {y}) $$

이제 y의 변화량 $ \dot {y(w)} $를 chain rule로 나타내고 위에서 구한 $\dot {w}$를 대입하면 다음과 같습니다.

$$ \dot{y(w)} = \triangledown y(w)^{T}\dot {w} = -\triangledown y(w)^{T} \triangledown y(w) (y(w)-\bar {y}) $$

위의 식에서 $ \triangledown y(w) $는 parameter에 대한 gradient이므로 non-linear입니다.

$\triangledown y(w)$를 kernel의 basis function $\phi(x)$라고 생각하면 kernel $ H(w)=\triangledown y(w)^{T} \triangledown y(w) $라고 치환해줄 수 있습니다. 이때 H(w)를 Neural Tangent Kernel(NTK)라고 합니다!! (이렇게 kernel로 치환하여 생각하면 non-linear 하여 풀지 못했던 식을 쉽게 만들어 주어 해가 존재함을 보일 때 중요하게 사용됩니다.) 또한 위에서 본 것처럼 width가 무한 대면 $ \triangledown y(w(t)) \approx \triangledown y(w_{0}) $이므로 $ H(w(t)) \approx H(w_{0}) $ 가 성립합니다. 정리해서 다시 쓰면 다음과 같습니다.

$$ \dot {y(w)} = -H(w_{0})(y(w)-\bar {y}) $$

$ y(w)-\bar {y} $를 $u$라고 치환하면 다음과 같습니다.

$$ \dot {u} = -H(w_{0})u $$

$$ u(t) = u(0) e^{-H(w_{0})t} $$

이때, $H(w_{0})$가 positive definite 하여(논문에서 증명) t가 증가할 때 $u(t)$는 항상 0으로 가게 됩니다!! $u(t)$는 prediction - label이었으므로 training loss가 0이 된다는 것을 의미합니다.

이렇게 Neural Tanget Kernel을 사용하여 width가 무한대일 때 loss를 0으로 만드는 해가 반드시 존재함을 보였습니다.

 

Reference

[1] https://arxiv.org/abs/1806.07572

Neural Tangent Kernel: Convergence and Generalization in Neural Networks

[2] https://rajatvd.github.io/NTK/

Understanding the Neural Tangent Kernel