Good news ! Spotting error in NTK's empirical testing code !

Today I’ve spotted an error in the code for empirically testing NTK’s along a huge variety of network configurations (Depth, Width, Dataset size, etc.).

In fact the parameterization of the NTK is not correct in the Neural-Tangents library. because of homogeneity of ReLU, the weights of the last layer are not correctly scaled.

I tried to see if the convegence appeared at the good scale (1/width) to spot where the error was, because even the Neural-Tangents documentation and repo code was not clear about Xavier’s initialization.

I’ve been working on this for the past 2 weeks, and I’m very happy to have found it ! The moment you found something stressing before sleeping, try to fix it 1 hour like i’ve done, and if it doesn’t work, make a wise choice and go to sleep ! It’s worth it every time !

I had long debugging sesh in my life, and it’s the first time I decide to allocate 1 hour to debug it, and it worked !