Linearize function explanation

[eeshaan-ravi-tivari]

Can someone explain exactly how does linearize() function works and what is it using as inputs, what & how is it calculating that it is supposed to return i.e. first-order Taylor approximation of the function? An example to answer all these in detail will b much helpful. Also if I want to replicate its behaviour in Pytorch, how can I do that?

[romanngg]

Correct, nt.linearize(f, p0) is the first-order Taylor approximation to f at p0, e.g. np.linearize(f, p0)(p) = f(p0) + [df/dp]_{p0} (p - p0). It's implemented as a simple call to jax.jvp

neural-tangents/neural_tangents/_src/empirical.py

Line 124 in 239cc84

def f_lin(p, *args, **kwargs):

and you can read more about it in
https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#jacobian-vector-products-jvps-aka-forward-mode-autodiff
or in papers about JAX
https://arxiv.org/abs/2105.09469
https://arxiv.org/abs/2204.10923

Torch has https://pytorch.org/docs/stable/generated/torch.autograd.functional.jvp.html, but please note their warning

autograd.functional.jvp computes the jvp by using the backward of the backward (sometimes called the double backwards trick). This is not the most performant way of computing the jvp. Please consider using [functorch’s jvp](https://github.com/pytorch/functorch#jvp) or the [low-level forward-mode AD API](https://pytorch.org/docs/stable/autograd.html#forward-mode-ad) instead.

so perhaps https://github.com/pytorch/functorch#jvp is a better entry point.

[eeshaan-ravi-tivari]

Actually what I want to do is, right now I have model build using pytorch framework (via usual nn.Module, optimizers etc). Now I want to use the 1st order taylor approximation of these model and train them using gradient descent (just as what is done in my other post's code using ntk.linearize() where you replied). So can you put in a right direction if this is possible or not and if yes how should it be done?

[romanngg]

Here's a quick rough example - https://colab.research.google.com/gist/icml2022anon/3ba1bb0fd4fedabd651ded60fafcf219/torch_linearize.ipynb please check for correctness, I'm a total noob in pytorch, and essentially just translated the code from

neural-tangents/neural_tangents/_src/empirical.py

Line 124 in 239cc84

def f_lin(p, *args, **kwargs):

into functorch.

Then you can either follow https://github.com/google/neural-tangents/blob/main/notebooks/weight_space_linearization.ipynb to train your f_lin, just use functorch.grad instead of jax.grad, or perhaps alternatively you can create a new nn.Module which will call f_lin in the forward pass; it will need to store the input params passed to f_lin as trainable parameters, and the linearization params passed to linearize as the non-trainable parameters. I'm not familiar with this API and how it interacts with torch optimizers, so you may know better how to proceed here.

[eeshaan-ravi-tivari]

what do u mean by - "it will need to store the input params passed to f_lin as trainable parameters, and the linearization params passed to linearize as the non-trainable parameters. "
?

[romanngg]

I just wanted to emphasize that the linearized model has two parameters: those around which you Taylor expand (in

neural-tangents/neural_tangents/_src/empirical.py

Line 100 in 239cc84

def linearize(f: ApplyFn, params: PyTree) -> ApplyFn:

it's params passed to linearize decorator along with f), and those at which you evaluate your linearized model (p in the link above that you pass to f_lin), and training a linearized models conventionally refers to training the parameters p at which you evaluate f_lin, while params stay frozen. So if you store p and params somewhere in the nn.Module, you perhaps need to label p as trainable and params as not, for torch optimizers to treat them properly. But I haven't looked at how nn.Module and torch optimizers work, so this comment may not be very useful.