학습의 끝에서
알고리즘이 모습을 드러낸다

책의 앞부분은 신경망을 넓게 만들수록 이론이 단순해진다는 사실을 이용했다. infinite width에서는 stochastic fluctuation이 사라지고 NTK가 frozen되며, fully-trained solution은 optimization algorithm의 세부 경로와 무관해졌다. 그러나 이 편리함에는 대가가 있었다. representation learning이 사라졌다.

이 장은 그 대가를 되돌려 받는다. 폭을 유한하게 두면 NTK는 더 이상 고정되지 않고, dNTK가 NTK의 움직임을 만들며, 그 dNTK 자체도 움직인다. 최종 함수는 단순한 kernel interpolation이 아니라 training trajectory의 흔적을 포함한다. 저자들이 이 장에서 해결하려는 문제는 명확하다. finite-width network를 끝까지 학습했을 때의 ensemble distribution을 초기화 시점의 통계량만으로 기술하는 것이다.

이를 위해 두 가지 현실을 받아들여야 한다. 첫째, \(O(1/n)\)까지 self-consistent하려면 Chapter 11의 dNTK만으로 부족하다. output update를 \(\eta^3\), NTK update를 \(\eta^2\), dNTK update를 \(\eta\)까지 밀어야 하며 두 종류의 ddNTK가 등장한다. 둘째, 이 Taylor 전략은 ReLU처럼 nonsmooth한 activation의 finite-width dynamics에서는 깨진다.

이후의 전개는 세 단계로 읽을 수 있다. 먼저 finite-width dynamics를 닫는 ddNTK statistics를 만든다. 그 다음 두 번의 Newton-like update와 many-step vanilla gradient descent를 각각 풀어 “학습의 끝”에 도달한다. 마지막으로 두 해를 비교해 optimizer의 inductive bias를 sample space의 algorithm projector로 분리한다.

Output update: 세 번째 order에서 첫 ddNTK가 나온다

parameter update 이후의 output을 세 번째 order까지 Taylor-expand하면 Chapter 11에서 본 NTK와 dNTK 뒤에 새로운 cubic term이 붙는다.

첫 번째 ddNTK, \(dd_I\widehat H\), 는 network output의 parameter 3차 derivative와 세 개의 1차 derivative를 learning-rate tensor로 contraction한 object다. 즉 quadratic feature motion 다음 단계의 curvature를 담는다. 이 항의 initialization statistics가 \(O(1/n)\)이기 때문에 버릴 수 없다.

NTK와 dNTK도 함께 움직여야 한다

output을 cubic order까지 썼다면 NTK 역시 quadratic order까지 업데이트해야 한다. 여기에는 기존 dNTK와 첫 ddNTK뿐 아니라 두 번째 ddNTK, \(dd_{II}\widehat H\), 가 등장한다. 이 두 번째 object는 두 개의 output Hessian과 두 개의 gradient를 조합한 구조다. dNTK update도 leading order에서 두 ddNTK에 의해 구동된다.

이 세 update equation은 \(O(1/n)\)에서 필요한 finite-width dynamical system을 닫는다. 더 높은 Taylor term은 이 order에서는 \(O(1/n^2)\)로 내려간다. 요컨대 Chapter 11의 “feature가 움직인다”는 정성적 설명이 여기서 “feature와 meta-feature가 함께 움직이는 coupled dynamics”로 완성된다.

R, S, T, U: 거대한 tensor를 네 개의 sample-space object로 압축한다

fully-trained distribution을 계산하려면 initialization에서

를 알아야 한다. ddNTK의 leading statistics는 mean에서 나오며 variance와 higher cross-correlation은 subleading이다. 저자들은 첫 ddNTK mean을 \(R\), 두 번째 ddNTK mean을 \(S,T,U\)라는 sample-index tensor로 분해한다. 네 object는 모두 \(O(1)\)이지만 전체 ddNTK mean 앞에 \(1/n\)이 붙는다.

차원 분석을 하면 \(R,S,T,U\)는 NTK의 세제곱에 해당하는 차원을 갖는다. 따라서 자연스러운 dimensionless ratio는 \(RK/(n\Theta^3)\), \(SK/(n\Theta^3)\), \(TK/(n\Theta^3)\), \(UK/(n\Theta^3)\)가 된다. 이 normalization은 뒤의 end-of-training prediction에서 그대로 다시 등장한다.

\(K^\star=0\) class에서는 R, S, T가 모두 \(L/n\)의 relevance를 가진다

odd activation, 특히 tanh 같은 \(K^\star=0\) class를 criticality와 learning-rate equivalence에 맞추면 single-input solution에서 \(p_R=p_S=p_T=-2\), \(p_U=0\)을 얻는다. normalized ratio로 보면 \(R,S,T\)는 모두 layer index를 따라 \(\ell/n\)으로 커지고, \(U\)는 \(1/n\)에 머물러 subdominant다.

중요한 점은 “higher derivative가 하나 더 생겼으니 더 작은 효과”가 아니라는 것이다. properly normalized하면 ddNTK도 dNTK와 NTK fluctuation과 같은 leading effective-theory cutoff에 참여한다.

이 장의 가장 중요한 제한은 숨기지 않는 편이 좋다. nonlinear scale-invariant activation, 대표적으로 ReLU는 smooth하지 않다. \(\sigma'\)는 step function이고 \(\sigma''\)는 Dirac delta이며, 더 높은 derivative는 delta derivative를 포함한다. ddNTK recursion에서 필요한 Gaussian expectation에 \(\sigma'''\)까지 들어가면 일부 적분은 발산하려 한다.

하지만 이 문제를 distribution calculus의 기술적 난점으로만 보면 핵심을 놓친다. 책은 bias 하나를 가진 단일 ReLU neuron으로 더 물리적인 설명을 준다.

update 전에는 \(x+b>0\)여서 neuron이 켜져 있다고 하자. 작은 update 뒤에도 부호가 유지되면 output change는 \(O(\eta)\)이고 Taylor expansion이 잘 작동한다. 그러나 update가 threshold를 넘어 \(x+b-\eta\varepsilon<0\)이 되면 activation은 갑자기 꺼진다. 이때 output change는 \(-x-b\), 즉 \(O(1)\)일 수 있다. \(\eta\)에 대한 analytic expansion으로는 threshold crossing을 잡을 수 없다.

원문은 deep ReLU MLP의 normalized NTK update를 실제로 측정하면 perturbative \(L/n\) 대신 width에 대해 대략 \(1/\sqrt n\), depth에 대해 linear한 scaling이 보인다고 말한다. 이는 이 chapter의 finite-width Taylor dynamics가 nonsmooth ReLU에 정량적으로 맞지 않는다는 직접적인 경고다.

따라서 분석은 smooth critical activation으로 돌아간다. 저자들은 GELU와 SWISH처럼 ReLU를 smooth하게 만든 activation 중 criticality를 허용하는 경우에는 이 방법을 적용할 수 있다고 본다. scale-invariant class 안에서 정확히 남는 특수한 경우는 deep linear network다. linear activation에서는 \(R=U=0\)이지만 두 번째 ddNTK의 \(S,T\)가 nonzero이고 역시 leading finite-width dynamics에 들어온다.

Fully trained라는 말의 수학적 뜻

training set \(A\)의 모든 sample에 대해 \(T\) step 뒤 output이 target과 같아지면 fully trained라고 정의한다.

그리고 test input을 포함한 일반 input의 final output을 initialization에서의 \(z,\widehat H,d\widehat H,dd_I\widehat H,dd_{II}\widehat H\)의 functional로 쓰는 것이 목표다. 이 functional을 얻으면 initialization joint distribution을 밀어 보내 fully-trained ensemble distribution을 계산할 수 있다.

Infinite-width giant leap는 finite width에서 O(1/n)만큼 빗나간다

Chapter 10의 Newton-like giant leap를 finite-width network에 그대로 적용해 보자. NTK mean의 inverse를 이용하면 leading order에서는 training target에 도착한다. 그러나 더 세밀하게 보면 residual error가 남는다.

이 \(1/n\) residual에는 두 종류가 섞여 있다. 하나는 model realization마다 다른 stochastic NTK fluctuation이다. 다른 하나는 step을 밟는 동안 feature가 실제로 변했다는 representation-learning correction, 즉 dNTK와 ddNTK가 만드는 nonlinear effect다. 이를 모두 보정하면 이론적 residual을 \(O(1/n^2)\)까지 내릴 수 있다.

첫 update: 각 network 자신의 stochastic NTK를 뒤집는다

finite width에서는 output component끼리 wiring된다. 따라서 sample-space Newton tensor에 output-component index \(i,j\)까지 붙인 generalized tensor \(\kappa_{ij;\tilde\alpha_1\tilde\alpha_2}\)가 필요하다. 첫 giant leap는 training set에서 각 model realization의 complete stochastic NTK inverse를 사용한다.

무한폭에서는 모든 model이 사실상 같은 NTK를 가졌지만, 이제는 각 realization에 맞는 Newton step을 따로 취한다. 이 step이 초기 \(O(1)\) training error를 \(O(1/n)\)으로 줄인다. 남은 것은 거의 전부 representation 변화가 남긴 correction이다.

두 번째 update: representation learning이 되돌아온다

첫 step 이후 residual 자체가 이미 \(O(1/n)\)이므로 두 번째 update는 작다. 이 단계에서는 dNTK와 ddNTK term을 다시 명시적으로 넣을 필요가 없다. 작은 residual에 mean NTK inverse를 곱해 correction하면

이 두-step algorithm은 practical optimizer 제안이 아니다. large training set에서 stochastic NTK tensor를 inversion하는 것은 비현실적이다. 저자들의 목적은 “왜 finite-width network는 한 번의 kernel solve로 끝나지 않는가”를 analytically 분해하는 데 있다.

이제 이론적 giant leap 대신 vanilla gradient descent를 many many steps 수행한다. MSE와 constant global learning rate \(\eta\) 아래에서 output, NTK, dNTK가 서로 coupled된 difference equation을 이룬다. ddNTK는 이 leading order에서는 initialization 값으로 고정해도 된다.

Free theory: dNTK와 ddNTK를 끄면 익숙한 exponential convergence가 나온다

먼저 nonlinear interaction을 끈다. 그러면 NTK는 initialization에서 고정되고 residual error는 선형 difference equation을 따른다.

training-set stochastic NTK가 positive definite이고 \(\eta\)가 충분히 작아 \(\|I-\eta\widehat H\|<1\)이면 \(U(t)\)는 exponentially zero로 간다. 따라서 training error도 exponentially 사라진다. free end-of-training prediction은 stochastic NTK를 이용한 kernel-like solution이며, dNTK와 ddNTK를 모두 껐을 때는 앞 절의 Newton-like solution과 일치한다.

이 일치는 아주 중요한 진단이다. finite width에서 optimizer dependence가 생기는 원인은 random NTK fluctuation 자체가 아니라 NTK differential이 만드는 nonlinear dynamics라는 뜻이다.

Interacting theory: free part 위에 finite-width interaction을 perturbation으로 얹는다

다음으로 output과 NTK를 free part와 interaction part로 나눈다.

\(z^I\)와 \(H^I\)는 이미 dNTK 또는 ddNTK에 proportional하므로, 이 correction과 다른 NTK differential을 다시 곱한 항은 higher order로 버릴 수 있다. 이 perturbative bookkeeping 덕분에 겉보기에는 nonlinear coupled system인 문제를 순차적으로 풀 수 있다.

dNTK도 학습한다: meta-representation learning

dNTK dynamics의 closed form을 보면 training error가 initialization에서 얼마나 줄었는지가 dNTK change를 결정한다. residual이 크게 줄수록 dNTK가 더 많이 바뀐다. Chapter 11의 parameter-space 언어로 번역하면 meta feature function 자체가 학습되는 것이다. 저자들은 이를 한 단계 높은 representation dynamics로 해석한다.

a(t), b(t), c(t): 긴 시간합을 세 개의 helper로 압축한다

many-step solution에는 free residual의 linear, quadratic, cubic time sum이 반복해서 나온다. 이를 각각 \(a(t),b(t),c(t)\)로 묶고, geometric sum의 inverse 역할을 하는 \(X_I,X_{II},X_{III}\) tensor를 정의한다. 복잡한 algebra의 핵심은 이 helper들이 “얼마나 오랫동안 얼마만큼의 training error가 남아 있었는가”를 order별로 적분한 memory라는 점이다.

interaction NTK는 \(a,b\)로, interaction output의 damping force는 \(a,b,c\)로 표현된다. training이 진행되면 free residual이 exponential하게 줄기 때문에 이 time sum도 수렴한다.

끝까지 가면 training error는 사라진다

같은 stability condition 아래 free residual과 damping force가 모두 exponential하게 decay하며, interaction part의 training-set output 역시 zero로 수렴한다.

다만 여기서 “0”은 effective theory의 truncation 안에서의 statement다. 이 장은 \(O(1/n)\), 더 정확히 deep scaling을 반영하면 leading \(O(L/n)\) theory이므로 누락된 error는 다음 order인 \(O((L/n)^2)\) 수준으로 이해해야 한다. convergence 조건 역시 positive-definite NTK와 충분히 작은 learning rate를 전제로 한다.

Fully-trained gradient-descent solution

긴 time sums를 모두 모으면 vanilla gradient descent로 fully-trained finite-width network의 closed-form solution, Eq. ∞.141에 도달한다. 식은 매우 길지만 구조는 짧게 읽을 수 있다.

여기서 여섯 개의 \(Z\) tensor가 algorithm projectors다. \(Z_A,Z_B\)는 dNTK correction, \(Z_{IA},Z_{IB}\)는 첫 ddNTK correction, \(Z_{IIA},Z_{IIB}\)는 두 번째 ddNTK correction에 training error를 어떻게 투영할지 정한다.

두-step Newton-like algorithm도 같은 skeleton을 갖지만 projector의 값이 다르다. gradient flow의 \(\eta\to0\) limit에서도 다시 다른 projector를 얻는다. 저자들은 MSE finite-width problem의 leading solution이 일반적으로 이 universal form을 가지며 optimizer 차이는 projector에만 모일 것이라고 conjecture한다.

예를 들어 책은 \(L=10\), \(n=100\)인 network라면 neglected next-order scale을 대략 \((L/n)^2=1\%\)로 본다. 이것은 empirical accuracy error가 1%라는 뜻이 아니다. effective-theory description의 truncation scale을 예시한 것이다.

Trained NTK로 다시 압축할 수 있다

Chapter 11의 nearly-kernel method처럼 finite-width prediction도 trained NTK를 정의하면 kernel-like form으로 다시 쓸 수 있다.

겉모양은 kernel prediction이지만 \(\overline H\)는 training-dependent다. 이 압축은 representation learning이 “kernel-like language를 버려야만” 표현되는 것이 아니라, kernel 자체를 learned object로 바꾸면 sample space에서도 표현할 수 있다는 뜻이다.

Finite-width prediction은 더 이상 Gaussian이 아니다

infinite width에서 final output은 Gaussian initial output의 linear functional이어서 Gaussian이었다. finite width에서는 prediction이 \(z,\Delta H,dH,dd_IH,dd_{II}H\)의 nonlinear functional이다. initialization joint distribution이 nearly Gaussian이어도 nonlinear mapping을 통과한 fully-trained distribution은 exactly Gaussian이 아니며 connected higher-point correlator가 생긴다.

Mean prediction에는 네 종류의 finite-width correction이 들어간다

ensemble mean은 infinite-width와 비슷한 NTK prediction 위에 네 correction family가 얹힌다: NTK fluctuation, dNTK, ddNTK-I, ddNTK-II다. 주목할 점은 NTK variance도 mean prediction을 바꾼다는 것이다. 각 network realization이 자기 고유의 stochastic NTK로 training data를 fit하기 때문이다.

또한 finite width에서는 output component wiring이 나타난다. training label의 \(j\)-th component가 test prediction의 \(i\)-th component에 영향을 줄 수 있으며, covariance에서도 서로 다른 output component가 coupling된다. 이는 infinite-width no-wiring result와 질적으로 다른 현상이다.

Generalization correction도 \(L/n\)이다

mean prediction의 finite-width contribution을 infinite-width kernel term과 비교하면 NTK variance의 \(A,B\), dNTK의 \(P,Q\), ddNTK의 \(R,S,T,U\)가 모두 dimensionless ratio로 나타난다. 앞 장들과 이 장에서 얻은 scaling law를 모으면 subdominant \(U\)를 제외한 leading ratio는 모두 \(L/n\)으로 간다.

variance 쪽의 leading correction도 같은 order로 예상된다. 따라서 \(r=L/n\)은 단순한 approximation parameter를 넘어, representation-learning benefit과 random-instantiation cost를 동시에 조절하는 축이 된다.

이 장만으로 최적 \(L/n\)을 계산할 수는 없다. leading correction이 \(r\)에 선형이므로 derivative를 취해도 optimum이 생기지 않는다. optimum을 찾으려면 \(O(L^2/n^2)\) higher-order term까지 필요하다. 이 장이 말할 수 있는 것은 “작은 nonzero \(r\)가 generalization을 개선하는지 악화하는지”를 coefficient의 부호로 판단하는 정도다.

Optimizer의 inductive bias를 architecture의 inductive bias와 분리한다

algorithm projector는 이 장의 가장 흥미로운 개념적 결과다. parameter space에서 optimizer는 weight update rule로 정의되지만, end of training에 가면 그 세부 경로는 final parameter 안에 숨어 버린다. projector는 이 숨은 경로를 sample space에서 다시 드러낸다.

더 중요한 것은 projector가 모든 finite-width effect를 마음대로 바꾸지는 못한다는 점이다. NTK fluctuation이 만드는 irreducible variance contribution은 training algorithm과 무관하다. 반면 \(P,Q,R,S,T,U\)로 표현되는 NTK-differential contribution, 즉 representation-learning part에는 projector가 직접 작용한다. 따라서 optimizer는 random fluctuation의 비용을 없앨 수는 없지만, 그 비용이 고정된 상태에서 representation-learning benefit을 어떤 방향으로 사용할지는 바꿀 수 있다.

저자들은 여기서 한 걸음 더 나아가 inverse algorithm design 가능성을 제안한다. 먼저 원하는 projector의 functional form을 설계하고, 그 projector를 만드는 parameter-space optimizer를 역으로 찾는 방식이다. 이는 이 책이 수행한 결과가 아니라 future direction에 가까운 제안이다.

Finite width에서는 gradient descent와 exact Bayes가 다시 갈라진다

Chapter 10에서는 특별한 hyperparameter 선택 아래 infinite-width gradient descent와 exact Bayesian inference가 일치하는 경우가 있었다. hidden-layer learning rate를 0으로 두면 마지막 layer만 학습하는 random-feature model이 되고 \(\Theta=K\)가 되어 두 해가 같아졌다.

finite width에서는 이 특수한 동일성이 유지되지 않는다. hidden-layer learning을 끄면 hidden-layer NTK가 사라지고, 그 결과 \(B,F,P,Q\)와 ddNTK effect까지 모두 꺼진다. gradient-based model은 representation learning을 잃는다. 반면 finite-width Bayesian inference에서는 representation learning correction이 남는다. 그러므로 posterior mean과 gradient-trained ensemble mean은 맞을 수 없다.

근본적인 차이는 필요한 통계 정보에도 드러난다. Bayesian inference는 preactivation distribution \(p(z)\)에 의존하지만 gradient-based learning의 finite-width solution은 \(z\)와 함께 \(H,dH,ddH\)처럼 model output의 parameter derivative statistics를 필요로 한다. 이 derivative 정보는 exact Bayesian posterior calculation에는 보이지 않는다.

§∞.3: 보기 싫은 식을 뒤로 숨긴 이유

마지막 절은 두 ddNTK의 full stochastic forward equations와 \(R,S,T,U\)의 complete multi-input layer recursion을 모아 둔다. 저자들은 이 식을 본문 흐름에서 떼어낸 이유를 세 가지로 정리한다.

• 01GD와 Bayes의 finite-width 불일치 확인. hidden-layer learning을 끈 hyperparameter에서 모든 NTK differential contribution이 실제로 사라지는지 full recursion으로 검증한다.
• 02depth-to-width scaling 검증. 간략한 single-input 결과가 full RG recursion과 일치하는지 확인한다.
• 03multi-input·higher statistics 계산 기반. 여러 input에서 ddNTK가 mean prediction과 higher connected correlator에 기여하는 항을 analytical 또는 numerical하게 평가할 수 있게 한다.

구조적으로는 \(dd_IH\)의 mean을 \(R\) 하나로, \(dd_{II}H\)의 mean을 \(S,T,U\) 세 tensor로 decomposition한 뒤 각각의 layer-to-layer recurrence를 제시한다. 모든 식을 읽는 것보다 중요한 것은 이 절이 “앞에서 사용한 scaling result가 어디서 왔는지”를 닫는 appendix-like 증명 장치라는 점이다.

이 장이 남기는 경계와 의미

• 01적용 범위. smooth activation을 전제로 한 \(\eta\)-Taylor finite-width dynamics이며 nonsmooth ReLU에는 정량적으로 적용되지 않는다.
• 02정확도 범위. leading finite-width theory는 \(O(L/n)\) correction을 담고 다음 neglected scale은 대략 \(O((L/n)^2)\)이다.
• 03generalization optimum. 최적 aspect ratio를 얻으려면 higher-order theory가 필요하다. 이 장은 leading tradeoff를 설명한다.
• 04optimizer theory. six-projector universal form은 MSE에 대해 explicit solution과 conjecture로 제시된다. cross-entropy에서는 유사하지만 다른 form이 필요하다고 저자들이 명시한다.
• 05가장 큰 질적 변화. infinite width의 free·linear·algorithm-independent regime에서 finite width의 interacting·nonlinear·algorithm-dependent regime으로 넘어간다.