A small-scale reproduction of the Neural Tangent Kernel phenomenon (Jacot, Gabriel & Hongler, 2018; Lee et al., 2019) on a 1D regression task, with an original activation-function comparison testing whether smoother activations give tighter linearizations at fixed width.

A two-layer MLP under NTK parameterization is trained on $y = \sin(\pi x)$ across widths $N \in {10, 100, 1000, 10000}$. At $N = 10000$ the trained network and its linearized surrogate are visually indistinguishable, and the network's predictions agree with closed-form kernel regression using the empirical NTK to within 0.005. The fluctuations of the empirical NTK across random initializations scale as $N^{-0.54}$, in close agreement with the theoretical $N^{-1/2}$ central-limit prediction.

• Phase 1 — Linearization phenomenon. At fixed initialization, the real network and its linearization make increasingly similar predictions as width grows. Mean test-set gap shrinks from 0.044 at $N = 10$ to 0.004 at $N = 10000$.
• Phase 2 — Empirical NTK. The empirical NTK at $\theta_0$ is well-formed and consistent with the closed-form arc-cosine kernel for ReLU. Across 20 random initializations at each width, NTK fluctuations scale as $N^{-0.54}$ (theory: $N^{-1/2}$). Closed-form kernel regression with the empirical NTK reproduces the trained network's predictions at large width, with informative deviations at small width where the network escapes the lazy regime.
• Phase 3 — Activation-function comparison. The physics-flavored prediction that smoother activations should give tighter linearizations is confirmed: across ${$ReLU, Tanh, GELU$}$ at every width, the ordering is monotone with activation smoothness, with GELU giving up to $4.7\times$ tighter linearization than ReLU. Smooth-activation slopes ($-0.40$, $-0.42$) are close to the theoretical $-1/2$; ReLU's slope ($-0.26$) is substantially shallower, consistent with non-smooth-activation corrections dominating at finite width.

A short report with all figures and full discussion is in report.pdf.

git clone https://github.com/QuantumPouya19/ntk-linearization.git
cd ntk-linearization
python -m venv .venv
# Windows:
.venv\Scripts\activate
# macOS/Linux:
source .venv/bin/activate
pip install -r requirements.txt

Then, in order:

# Phase 1: train across widths, produce the linearization-overlap figure.
python scripts/plot_width_convergence.py

# Phase 2: empirical NTK and the 1/sqrt(N) scaling.
python scripts/plot_ntk_heatmap.py
python scripts/plot_ntk_fluctuations.py     # ~10 minutes
python scripts/plot_kernel_regression.py

# Phase 3: activation-function comparison.
python scripts/run_activation_comparison.py

All hyperparameters live in configs/default.yaml. Runs are seeded; the default seed is 0. The total CPU runtime is around 15 minutes, dominated by Phase 2's fluctuation analysis at width 10000.

ntk-linearization/
├── src/                                 # Library code
│   ├── config.py                        # Dataclass + YAML loader
│   ├── data.py                          # 1D regression dataset
│   ├── model.py                         # NTK-parameterized MLP
│   ├── train.py                         # Training loop
│   ├── linearize.py                     # Linearized surrogate via Jacobian
│   ├── kernel.py                        # Empirical NTK + kernel regression
│   └── utils.py                         # Seeding
├── scripts/                             # One-purpose entry points
│   ├── plot_width_convergence.py
│   ├── plot_ntk_heatmap.py
│   ├── plot_ntk_fluctuations.py
│   ├── plot_kernel_regression.py
│   └── run_activation_comparison.py
├── configs/
│   ├── default.yaml                     # Phase 1 & 2 hyperparameters
│   └── activations.yaml                 # Phase 3 sweep config
├── results/                             # Figures and JSON outputs
├── notes/                               # Written discussion per phase
├── report.pdf                           # Full writeup
├── requirements.txt
├── LICENSE
└── README.md

Several setup details that the canonical writeups underspecify are worth flagging:

• NTK parameterization with explicit $1/\sqrt{N}$ in the forward pass is essential for the theorem to apply as stated. Using standard PyTorch Kaiming initialization without the forward-pass rescaling does not give the NTK limit.
• The first-layer bias must be initialized from a nontrivial distribution. A zero-initialized bias produces a degenerate two-dimensional feature space on 1D positive inputs, and the linearization cannot escape it during training.
• Full-batch SGD with learning rate $\sim 0.1$ is the right regime under this parameterization. The Adam optimizer adapts per-parameter learning rates and breaks the linearization story.
• Target functions must be representable in the NTK feature space. High-frequency targets fail to be fit due to spectral bias, even at large width.

A separate set of failed runs during development confirmed that getting any of these wrong produces silent failures that look superficially like the right setup.

This is a portfolio practice project. The NTK theorem (Jacot et al., 2018) and the linearization picture (Lee et al., 2019) are not new; the activation-function dependence of linearization quality is implicit in the standard Taylor analysis. The goal is a clean, reproducible worked example at the smallest non-trivial scale.

• Jacot, A., Gabriel, F., & Hongler, C. (2018). Neural Tangent Kernel: Convergence and Generalization in Neural Networks. arXiv:1806.07572
• Lee, J., et al. (2019). Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent. arXiv:1902.06720
• Cho, Y., & Saul, L. K. (2009). Kernel Methods for Deep Learning. NeurIPS 2009.

MIT. See LICENSE.