Question-ResnetX implementation

[eeshaan-ravi-tivari]

Can someone help me with how can I implement say ResNet20 using the Neural tangents library?
Also, how can we use a custom activation function for the network in NTK?

[romanngg]

1. I don't have a specifically resnet20 implemented, but I imagine this example should give the general idea:
   https://github.com/google/neural-tangents#infinitely-wideresnet

AFAIK we have all the necessary building blocks, except for BatchNorm (we have stax.LayerNorm though).

2. Re custom activation function, please see these options:

• https://neural-tangents.readthedocs.io/en/latest/_autosummary/neural_tangents.stax.Elementwise.html
• https://neural-tangents.readthedocs.io/en/latest/_autosummary/neural_tangents.stax.ElementwiseNumerical.html

We also provide some basic customization options in many nonlinearities https://neural-tangents.readthedocs.io/en/latest/stax.html#elementwise-nonlinear such as slopes of rectified units, Hermite polynomial degree, shifts and scaling, etc.

Lmk if this helps!

[eeshaan-ravi-tivari]

using your reply I tried nt.batch(nt.empirical_ntk_fn(f, vmap_axes=0), batch_size=100, device_count=0) followed by
g_dd = kernel_fn(train['image'], None, params) // 1
g_td = kernel_fn(test['image'], train['image'], params) // 2
predictor = nt.predict.gradient_descent(g_dd, y_train) // 3
but it is taking too long to complete the execution // 1(the g_dd step). Why is it so and is there any way to reduce this excution time like reducing the training data maybe?

[romanngg]

Yes, compute cost is quadratic in the batch size of train['image'], so reducing it will speed things up; similarly, the cost is linear in the size of params, so trying a smaller model would also make it faster. Also, if you have multiple GPUs, you can parallelize it by using device_count=-1. Finally, check out #68 (comment) to see if you need to compute the full (batch_size x n_logits) x (batch_size x n_logits) NTK (in which case you'd need to set trace_axes=()), or perhaps you can assume that your NTK is constant-block diagonal, with n_logits identical batch_size x batch_size diagonal blocks; in the latter case, you can replace f with lambda params, x: f(params, x)[:, 0], or lambda params, x: jnp.sum(f(params, x), axis=1) / O**0.5, where O is your number of logits; this way you'll save your compute by a factor of O^2. See also https://openreview.net/pdf?id=ym68T6OoO6L (note that several papers using the empirical NTK have employed the above trick of shaving off the O^2 factor -- see section E in the paper) for details on the cost of computing the empirical NTK as a function of batch size / number of logits / parameters.

[eeshaan-ravi-tivari]

can you tell that does the following code reflect or represents any learning/deductions/algorithm from the NTK learning theory -

######################################
init_fn, f, kernel_fn = l_layer_FFNN(l=l,input_shape= train['image'].shape[1],
hidden_shape=200, output_shape=10)

_, params = init_fn(key, train['image'].shape)

learning_rate = 0.01
training_steps = 1000

opt_init, opt_update, get_params = optimizers.adam(learning_rate)
opt_init = jit(opt_init)

f_lin = nt.linearize(f, params)
lin_state = opt_init(params)

loss = lambda fx, y_hat: -np.mean(logsoftmax(fx) * y_hat)
grad_lin_loss = jit(grad(lambda params, x, y: loss(f_lin(params, x), y)))

f_lin = jit(f_lin)

################
# Training #
################

print ('Time\tLinear Loss\t Lin Accuracy')
print_every=100

X, Y = train['image'], train['label']
Y_true = np.argmax(Y,axis=1)

for i in range(training_steps):

lin_params = get_params(lin_state)
lin_state = opt_update(i, grad_lin_loss(lin_params, X, Y), lin_state)

if i % print_every == 0:
  out_lin = f_lin(lin_params, X)
  exact_loss_lin = loss(out_lin, Y)

  Y_pred_lin = np.argmax(out_lin,axis=1)

  acc_score = accuracy_score(Y_true,Y_pred_lin)

  print('{}\t{:.3f}\t\t{:.3f}'.format(i, exact_loss_lin,100*acc_score))

if acc_score == 1.0:
  break  

[romanngg]

IIUC the code above just demonstrates that you can train a linearized network with am optimizer like the non-linear network; you can compare the weights and predictions of the linearized (f_lin) and the non-linear (f) network, and see how much they agree; NTK theory says that for very wide networks, with a small learning rate, gradient descent will produce very similar results for f and f_lin (and exactly matching in the infinite width limit). Not sure if this answers your question, lmk if you meant something else!

[eeshaan-ravi-tivari]

Thanks for your great reply, it just clears what I wanted