DL Published: TBD

Neural Tangent Kernel, cz. 1: jądra, parametryzacja i problem niewypukłości

Abstrakt

Pierwsza część opracowania pracy Jacota, Honglera i Gabriela "Neural Tangent Kernel: Convergence and Generalization in Neural Networks"[1]. Wprowadzamy funkcje jądrowe i kernel trick na przykładzie jądrowej regresji grzbietowej, definiujemy w pełni połączoną sieć neuronową w parametryzacji stosowanej w pracy oraz pokazujemy na prostym przykładzie, że wypukły funkcjonał kosztu staje się niewypukły po złożeniu z funkcją realizacji sieci. Ta obserwacja motywuje centralny pomysł pracy: analizę uczenia w przestrzeni funkcji zamiast w przestrzeni parametrów.

Wstęp

Jacot i in. pokazują, że w granicy nieskończonej szerokości warstw ukrytych uczenie sieci neuronowej spadkiem gradientu można opisać jako spadek wzdłuż gradientu jądrowego, wyznaczonego przez obiekt nazwany jądrem stycznym sieci neuronowej (Neural Tangent Kernel, NTK). Żeby to stwierdzenie miało precyzyjny sens, potrzebny jest pewien kontekst: czym jest funkcja jądrowa, jak metody jądrowe korzystają z iloczynów skalarnych, w jakiej parametryzacji rozważana jest sieć i dlaczego w ogóle warto przenosić analizę do przestrzeni funkcji.

Ta część skupia się więc na samym wprowadzeniu do formalizmu na którym bazuje NTK. Nie definiujemy w niej jeszcze samego NTK ani nie omawiamy wyników dotyczących zbieżności, jest to temat planowanej części 2.

Formalizm Matematyczny

Czym jest funkcja jądrowa

Pojęcie "jądro" bywa mylące, bo w matematyce oznacza dwa zupełnie inne obiekty. W algebrze liniowej jądrem przekształcenia liniowego $ L: V \rightarrow W $ nazywamy zbiór wektorów zerowanych przez to przekształcenie: $ \ker{(L)} = \{ v \in V : L(v) = 0 \} $. Z kolei, w uczeniu maszynowym odnosimy się do funkcji jądrowej (kernel function), która jest formalnie definiowana jako Positive-definite kernel[2]. W dalszej części artykułu, mówiąc "jądro" lub "funkcja jądrowa", mamy na myśli wyłącznie obiekt w kontekście uczenia maszynowego.

Formalnie, funkcja jądrowa to przekształcenie $ k: X \times X \rightarrow \mathbb{R} $, które mierzy "podobieństwo" między dwoma punktami. Wymagamy od niej dwóch własności:

symetria $k(x, x') = k(x', x) \quad \forall x, x' \in X$
dodatnia półokreśloność $\sum_{i=1}^n\sum_{j=1}^n c_i c_j \ k(x_i, x_j) \geq 0 \quad \forall n \in \mathbb{N}, \ \forall x_1, \dots, x_n \in X, \ \forall c_1, \dots, c_n \in \mathbb{R}$

Kluczowy fakt, wynikający z twierdzenia Moore'a–Aronszajna, mówi, że funkcja $ k $ spełnia oba warunki wtedy i tylko wtedy, gdy istnieje przestrzeń Hilberta $ H $ i przekształcenie $ \phi: X \rightarrow H $ takie, że

k(x, x') = \langle \phi(x), \phi(x')\rangle_H \quad \forall x, x' \in X

gdzie $ \langle \cdot, \cdot \rangle_H $ oznacza iloczyn skalarny w przestrzeni $ H $. Innymi słowy, każde takie jądro jest iloczynem skalarnym w pewnej przestrzeni cech, nawet jeśli przekształcenia $ \phi $ nie znamy jawnie. Te dwa warunki bywają w literaturze uczenia maszynowego nazywane warunkiem Mercera, choć samo twierdzenie Mercera dotyczy czegoś węższego: rozkładu spektralnego ciągłego jądra na zwartej dziedzinie.

Kernel Trick

Naturalnie narzuca się pytanie: po co w ogóle wprowadzać nowy, dosyć restrykcyjny obiekt, skoro mamy już klasyczny iloczyn skalarny? Zastanówmy się, jak metody uczenia maszynowego korzystają z iloczynu. Dla przykładu skorzystamy z regresji liniowej.

Rozważmy zbiór danych $ \{(x_i, y_i) \}_{i=1}^N $, $ x_i \in \mathbb{R}^d $, $ y_i \in \mathbb{R} $. Szukamy wag $ w \in \mathbb{R}^d $ minimalizujących błąd kwadratowy z regularyzacją grzbietową ($ \lambda > 0 $):

w^* = \arg\min_w \sum_{i=1}^N(w^T x_i - y_i)^2 + \lambda\| w \|^2 = \arg\min_w \| Xw - y \|^2 + \lambda \|w \|^2

gdzie $ X \in \mathbb{R}^{N \times d} $ to macierz danych, a i-ty wiersz to $ x_i^T $. Rozwiązanie możemy zapisać w postaci zamkniętej:

w^* = (X^TX + \lambda I)^{-1}X^Ty

Postać ta wymaga odwrócenia macierzy $ X^T X + \lambda I \in \mathbb{R}^{d \times d} $, więc jej koszt zależy od wymiaru cech $ d $. Dane występują w niej przez same cechy $ x_i $, a nie przez relacje między parami punktów.

Możemy jednak wykazać, że optymalne $ w^* $ zawsze można zapisać jako kombinację liniową punktów danych (representer theorem):

w^* = \sum_{i=1}^N \alpha_i x_i = X^T \alpha

dla pewnego wektora współczynników $ \alpha \in \mathbb{R}^N $. Podstawiając $ w = X^T\alpha $ do minimalizowanej funkcji, otrzymujemy

\min_\alpha \| XX^T\alpha - y \|^2 + \lambda \| X^T \alpha \|^2

Rozwiązując względem $ \alpha $ otrzymujemy rozwiązanie w formie dualnej:

\alpha^* = (XX^T + \lambda I)^{-1}y

Jeśli macierz $ XX^T $ jest osobliwa, minimalizacja po $ \alpha $ ma więcej rozwiązań, ale wszystkie dają to samo $ w^* $.

Różnica między powyższymi dwoma formami jest istotna. W formie pierwotnej pracujemy na macierzy $ X^TX \in \mathbb{R}^{d \times d} $ zbudowanej z cech, natomiast w formie dualnej cała konstrukcja opiera się na macierzy $ XX^T \in \mathbb{R}^{N \times N} $, której elementy $ (i, j) $ to iloczyny skalarne $ x_i \cdot x_j $.

Mając wyliczone $ \alpha^* $, predykcję dla nowego punktu $ x $ liczymy jako:

\hat{y}(x) = {w^*}^Tx = (X^T\alpha^*)^Tx = {\alpha^*}^TXx = \sum_{i=1}^N \alpha_i^*(x_i \cdot x)

Zarówno przy wyznaczaniu $ \alpha^* $, jak i przy samej predykcji, dane wejściowe pojawiają się wyłącznie jako iloczyny skalarne między parami punktów. To tutaj następuje cała zmiana konstrukcji, bo skoro udało się przedstawić model w takiej postaci, to możemy każdy taki iloczyn zastąpić funkcją jądrową:

x_i \cdot x_j \rightarrow k(x_i, x_j)

Definiując macierz Grama $ K_{ij} := k(x_i, x_j) $, otrzymujemy jądrową regresję grzbietową:

\alpha^* = (K + \lambda I)^{-1} y, \quad \hat{y}(x) = \sum_{i=1}^N \alpha_i^* k(x_i, x)

Model zachowuje się tak, jakby trenował regresję liniową w przestrzeni cech $ H $ (na wektorach $ \phi(x_i) $), mimo, że $ \phi $ nigdy nie zostało jawnie obliczone.

Kernel Ridge Regression z jądrem RBF $ k(x, x') = \exp\left(-\frac{(x-x')^2}{2\sigma^2}\right) $ na zaszumionych danych $ y = \sin(2x) + \varepsilon $, zestawiona ze zwykłą regresją liniową bez jądra (ta sama regularyzacja $ \lambda $). Suwaki $ \sigma $ i $ \lambda $ przeliczają dopasowanie KRR na żywo.

σ (szerokość jądra) 0.40

λ (regularyzacja) 0.050

Definicja sieci neuronowej

Rozważamy sztuczną sieć neuronową w pełni połączoną, z warstwami oznaczonymi od $ 0 $ (wejście) do $ L $ (wyjście), z której każda zawiera $ n_0, \dots, n_L $ neuronów. Warstwy ukryte stosują nieliniowość $ \sigma: \mathbb{R} \rightarrow \mathbb{R} $, która spełnia warunek Lipschitza, jest dwukrotnie różniczkowalna i ma ograniczoną drugą pochodną. Na warstwie wyjściowej nie stosujemy aktywacji.

Sieć w powyższym sformułowaniu nazywamy funkcją realizacji sieci (ANN realization function) $ F^{(L)} : \mathbb{R}^P \rightarrow \mathcal{F} $. Pojęciem tym będziemy się posługiwać w dalszej części artykułu. Nazwa może sugerować jakiś zewnętrzny mechanizm "wykonujący" sieć. Faktycznie $ F^{(L)} $ to po prostu sama sieć traktowana jako mapowanie z przestrzeni parametrów do przestrzeni funkcji: każdemu wektorowi parametrów $ \theta $ przypisuje ona konkretną funkcję $ f_\theta $, jaką sieć przyjmuje przy tych parametrach. Ważne jest odróżnienie dwóch etapów tego procesu:

\underbrace{\theta}_{\text{parametry}} \xrightarrow{F^{(L)}} \underbrace{f_\theta}_{\text{cała sieć jako funkcja}} \xrightarrow{\text{ewaluacja w } x} \underbrace{f_\theta(x)}_{\text{konkretne wyjście}}

$ F^{(L)} $ zatrzymuje się na pierwszym etapie, zwraca całą funkcję, nie liczbę czy wektor. Dopiero podstawienie zmiennej $ x $ jako wejście sieci daje konkretne wyjście.

Wymiar przestrzeni parametrów $ P $ jest liczony ze wzoru:

P = \sum_{l=0}^{L - 1}(n_l + 1)n_{l + 1}

Dowód

Wybrana warstwa $ l $ ma $ n_l \cdot n_{l + 1} $ wag i $ n_{l + 1} $ parametrów bias, jeden dla każdego neuronu wyjściowego. Oznacza to, że jedna warstwa ma $ n_l \cdot n_{l + 1} + n_{l + 1} $ parametrów, czyli po uproszczeniu $ (n_l + 1)n_{l + 1} $. Sumując po każdej warstwie, dostajemy wzór na wymiar $ P $.

Parametry sieci są inicjalizowane jako zmienne niezależne o identycznym rozkładzie Gaussa $ \mathcal{N}(0, 1) $.

Dla ustalonego rozkładu $ p^{\text{in}} $ na przestrzeni wejściowej $ \mathbb{R}^{n_0} $, przestrzeń funkcji $ \mathcal{F} $ jest zdefiniowana jako $ \{ f: \mathbb{R}^{n_0} \rightarrow \mathbb{R}^{n_L} \} $. W tej przestrzeni, rozważamy seminormę $ \| \cdot \|_{p^{\text{in}}} $ zdefiniowaną za pomocą formy dwuliniowej

\langle f, g \rangle_{p^{\text{in}}} = \mathbb{E}_{x \sim p^{\text{in}}}[f(x)^T g(x)]

poprzez $ \| f \|_{p^{\text{in}}} := \sqrt{\langle f,f \rangle_{p^{\text{in}}}} $. Seminorma to uogólnienie normy, które nie wymaga warunku: jeżeli $ x \in X $ spełnia $ ||x|| = 0 $, wtedy $ x = 0 $.

Niech $ x_1, \dots, x_N \in \mathbb{R}^{n_0} $ oznacza zbiór punktów danych wejściowych. Definiujemy na nich rozkład empiryczny

p^{\text{in}}(x) = \frac{1}{N}\sum_{i=1}^N \delta_{x_i}(x)

gdzie $ \delta_{x_i} $ to miara Diraca skupiona w punkcie $ x_i $. Przy takim wyborze $ \| f \|_{p^{\text{in}}} $ zależy wyłącznie od wartości $ f $ w punktach $ x_1, \dots, x_N $: niezerowa funkcja, która zeruje się na wszystkich punktach danych, ma seminormę równą $ 0 $. Właśnie dlatego jest to seminorma, a nie norma.

Zapis $ x \sim p^{\text{in}} $ oznacza, że zmienna losowa $ x $ ma rozkład $ p^{\text{in}} $.

Jako konkretny przykład rozważmy sieć z artykułu o reprezentacjach grup cyklicznych w sieciach neuronowych, trenowaną na zadaniu dodawania modulo $ p $. Każda liczba $ a, b \in \mathbb{Z}_p $ kodowana jest jako wektor one-hot, a wejście sieci to ich konkatenacja $ x = [e_a; e_b] \in \mathbb{R}^{2p} $, gdzie $ e_b $ zajmuje przesunięty blok indeksów $ [p, 2p) $. Pełny zbiór treningowy składa się ze wszystkich par $ (a, b) \in \mathbb{Z}_p \times \mathbb{Z}_p $, czyli $ N = p^2 $ przykładów.

Wszystkie możliwe kodowania one-hot dla $ p = 7 $: blok $ a $ na indeksach $ [0, 7) $, blok $ b $ przesunięty na $ [7, 14) $.

Rozkład empiryczny $ p^{\text{in}} $ na zbiorze wszystkich par $ (a, b) \in \mathbb{Z}_7 \times \mathbb{Z}_7 $ ($ N = 49 $): masa $ 1/49 $ na każdym punkcie $ \delta_{x_i} $.

Rozkład jest jednostajny, ponieważ każda para $ (a, b) $ pojawia się w zbiorze treningowym dokładnie raz, więc rozkład empiryczny z definicji przypisuje każdemu punktowi tę samą masę $ 1/p^2 $. Marginalnie każda pojedyncza liczba pojawia się w zbiorze tyle samo razy: $ p $ razy jako $ a $ i $ p $ razy jako $ b $.

Wyjście sieci definiujemy jako

f_\theta(x) := \tilde\alpha^{(L)}(x;\theta)

czyli preaktywację ostatniej warstwy, bez zastosowania funkcji $ \sigma $ na końcu.

Preaktywacje $ \tilde\alpha^{(l)}(\cdot;\theta):\mathbb{R}^{n_0}\to\mathbb{R}^{n_l} $ i aktywacje $ \alpha^{(l)}(\cdot;\theta):\mathbb{R}^{n_0}\to\mathbb{R}^{n_l} $ definiujemy rekurencyjnie:

\alpha^{(0)}(x;\theta) = x

\tilde\alpha^{(l+1)}(x;\theta) = \frac{1}{\sqrt{n_l}}W^{(l)}\alpha^{(l)}(x;\theta) + \beta b^{(l)}, \qquad l=0,\dots,L-1

\alpha^{(l)}(x;\theta) = \sigma\big(\tilde\alpha^{(l)}(x;\theta)\big), \qquad l=1,\dots,L-1

gdzie $ W^{(l)} \in \mathbb{R}^{n_{l+1} \times n_l} $ to macierz wag, $ b^{(l)} \in \mathbb{R}^{n_{l+1}} $ wektor biasów, a $ \beta > 0 $ to stały, nietrenowany współczynnik, którym autorzy regulują wpływ biasów na uczenie.

W powyższym sformułowaniu sieci ważna jest zmiana podejścia do inicjalizacji parametrów. W tzw. standardowej parametryzacji czynnika $ \frac{1}{\sqrt{n_l}} $ nie ma we wzorze propagacji. Zamiast tego wagi losuje się z rozkładu o przeskalowanej wariancji, np. $ W_{ij}^{(l)}\sim\mathcal{N}(0,\frac{1}{n_l}) $ (inicjalizacja LeCuna; inicjalizacja He używa wariancji $ \frac{2}{n_l} $). Dla przykładu domyślna inicjalizacja warstwy nn.Linear w PyTorch losuje zarówno wagi, jak i biasy z rozkładu jednostajnego na przedziale $ \left[-\frac{1}{\sqrt{n_l}}, \frac{1}{\sqrt{n_l}}\right] $. W parametryzacji NTK skalowanie przeniesione jest z rozkładu do wzoru propagacji, a wszystkie parametry mają rozkład $ \mathcal{N}(0,1) $. O ile zbiór funkcji $ F^{(L)}(\mathbb{R}^P) $ jest taki sam dla obu parametryzacji, to pochodne funkcji realizacji po parametrach już się różnią. Ta różnica będzie istotna w części 2.

Motywacja: od przestrzeni parametrów do przestrzeni funkcji

Proces trenowania ANN polega na optymalizacji $ f_\theta $ w przestrzeni funkcji $ \mathcal{F} $ względem funkcjonału kosztu $ C: \mathcal{F} \rightarrow \mathbb{R} $, takiego jak np. błąd średniokwadratowy albo entropia krzyżowa. Wracając do opisu w postaci etapów, możemy przedstawić to następująco:

\underbrace{\theta}_{\text{parametry}} \xrightarrow{F^{(L)}} \underbrace{f_\theta}_{\text{cała sieć jako funkcja}} \xrightarrow{\text{ewaluacja w } x} \underbrace{f_\theta(x)}_{\text{konkretne wyjście}} \xrightarrow{C} \underbrace{C(f_\theta)}_{\text{koszt (liczba)}}

Nawet dla wypukłego funkcjonału kosztu $ C $, złożony koszt $ C \circ F^{(L)}: \mathbb{R}^P \rightarrow \mathbb{R} $ jest mocno niewypukły.

Przykład

Rozważmy sieć z jednym neuronem ukrytym i rzeczywistą funkcją aktywacji $ \tanh $: $ f_\theta(x) = \theta_2 \tanh(\theta_1 x) $, z dwoma skalarnymi parametrami $ \theta = (\theta_1, \theta_2) $ (waga wejściowa i wyjściowa, bez biasów). Sieć ma dopasować się do małego zbioru danych $ \{(x_i, y_i)\}_{i=1}^N $, a funkcjonał kosztu to standardowy błąd średniokwadratowy (MSE):

C(f) = \frac{1}{N}\sum_{i=1}^N \big(f(x_i) - y_i\big)^2

$ C $ jest wypukłym funkcjonałem na przestrzeni funkcji $ \mathcal{F} $ (to suma wypukłych parabol). Po złożeniu z siecią $ F^{(2)} $ otrzymujemy jednak niewypukłą powierzchnię kosztu w przestrzeni parametrów:

(C \circ F^{(2)})(\theta_1, \theta_2) = \frac{1}{N}\sum_{i=1}^N \big(\theta_2 \tanh(\theta_1 x_i) - y_i\big)^2

Punkt $ \theta_1 = \theta_2 = 0 $ jest punktem krytycznym niezależnie od danych (bo $ \tanh(0) = 0 $ zeruje zależność od $ \theta_2 $), a hesjan w tym punkcie wynosi

\nabla^2(C \circ F^{(2)})(0, 0) = \begin{pmatrix} 0 & H \\ H & 0 \end{pmatrix}, \qquad H = -\frac{2}{N}\sum_{i=1}^N x_i y_i

Wartości własne tej macierzy to $ \pm |H| $, więc jest ona nieokreślona zawsze, gdy $ \sum_i x_i y_i \neq 0 $, co dla niemal każdego zbioru danych jest prawdą. Punkt $ (0, 0) $ jest więc punktem siodłowym, a $ C \circ F^{(2)} $ niewypukła, mimo że sam funkcjonał $ C $ pozostaje wypukły.

Poniżej: $ 5 $ punktów danych wygenerowanych z $ y = 0{,}9\tanh(0{,}8x) $. Lewy wykres pokazuje aktualną funkcję sieci $ f_\theta $ na tle tych punktów, a prawy, krajobraz kosztu $ C \circ F^{(2)} $ w przestrzeni $ (\theta_1, \theta_2) $, z zaznaczonym aktualnym punktem $ \theta $. Suwaki przesuwają ten sam punkt jednocześnie na obu wykresach.

$ f_\theta(x) = \theta_2\tanh(\theta_1 x) $: funkcja sieci (niebieska krzywa) na tle danych treningowych (czerwone punkty).

Krajobraz kosztu $ C \circ F^{(2)} $ w przestrzeni $ (\theta_1, \theta_2) $: biały punkt to aktualne $ \theta $.

θ₁ (waga 1) 0.90

θ₂ (waga 2) 1.05

Aby ominąć trudności wynikające z optymalizacji w niewypukłej przestrzeni parametrów, autorzy proponują analizę dynamiki bezpośrednio w przestrzeni funkcji $ \mathcal{F} $. Wykazują, że podczas uczenia funkcja sieci $ f_\theta $ ewoluuje zgodnie ze spadkiem wzdłuż gradientu jądrowego (ang. kernel gradient) wyznaczonego przez jądro styczne sieci neuronowej (Neural Tangent Kernel, NTK). Podejście to pozwala badać proces zbieżności w przestrzeni $ \mathcal{F} $, w której sam funkcjonał kosztu $ C $ pozostaje wypukły.

W części 2 zdefiniujemy gradient jądrowy i samo NTK, a następnie omówimy główne wyniki pracy: w granicy nieskończonej szerokości warstw ukrytych NTK przy inicjalizacji zbiega do deterministycznego jądra granicznego i pozostaje stałe w trakcie uczenia, a dla kosztu średniokwadratowego pozwala to opisać ewolucję funkcji sieci w postaci zamkniętej.

References

Jacot, A., Hongler, C., & Gabriel, F. (2018). Neural Tangent Kernel: Convergence and Generalization in Neural Networks. NeurIPS.
Schölkopf, B., & Smola, A. J. (2002). Learning with Kernels: Support Vector Machines, Regularization, Optimization, and Beyond. MIT Press.