How to estimate the gradient of an argmin loss

medvisioner: <p>Suppose we have a neural network f_\theta(x), where x is the input and \theta is the network’s parameters.</p>
<p>For each \theta, we can minimize f_\theta(x) w.r.t. x and obtain the minimum point x*(\theta):=argmin_x f_\theta(x).</p>
<p>My question is: how to estimate the gradient: dx*(\theta) / d\theta ?</p>
<p>I did some surveys and found this gradient (in fact any gradient) can be estimated by the <a href="https://en.wikipedia.org/wiki/Finite_difference" rel="noopener nofollow ugc">finite difference</a> technique. However, it is rather slow when $\theta$ includes many parameters, e.g., f_\theta is a large neural network.</p>
<p>Any possible approach is welcomed.</p>

KFrank: <p>Hi Med!</p>
<aside class="quote no-group" data-username="medvisioner" data-post="1" data-topic="173060" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/45deac/48.png" class="avatar"> medvisioner:</div>
<blockquote>
<p>how to estimate the gradient: dx*(\theta) / d\theta ?</p>
</blockquote>
</aside>
<p>Use calculus to expand <code>f</code> to second order around <code>x_star</code> and <code>theta</code>.   Solve<br>
for <code>x = argmin f (x)</code> as a function of <code>theta</code> in the context of this second-order<br>
expansion and then compute the gradient of this <code>x</code> with respect to <code>theta</code>.  This<br>
will give you an analytic result for the desired gradient in terms of second-order<br>
partial derivatives of <code>f (x; theta)</code> with respect to <code>x</code> and <code>theta</code>.</p>
<p>Then use <code>autograd.grad()</code> twice in a row to compute numerically the required<br>
partial derivatives for your specific network, <code>f</code>, and compute from them the gradient<br>
of <code>x_star</code> with respect to the parameters of your network, the <code>theta</code>.</p>
<p><code>f (x_star + delta_x; theta + delta_theta)</code> is given, to second order, by:</p>
<pre><code class="lang-plaintext">f (x_star; theta)  +
    delta_x * d_f/d_x  +  delta_theta * d_f/d_theta  +
    1/2 * delta_x^2 * d^2_f/d_x^2  +  delta_x * delta_theta * d^2_f/d_x d_theta + 1/2 delta_theta^2 * d^2_f/d_theta^2
</code></pre>
<p>where for the <code>theta</code> terms <code>*</code> is understood to denote the appropriate<br>
matrix-vector multiplication.</p>
<p>Keeping only the terms that depend on <code>x</code> and recognizing that <code>d_f/d_x</code> is zero<br>
at <code>x_star</code> (the minimum of <code>f (x)</code>), we get</p>
<pre><code class="lang-plaintext">f (x_star + delta_x; theta + delta_theta) ~
    1/2 * delta_x^2 * d^2_f/d_x^2  +  delta_x * delta_theta * d^2_f/d_x d_theta
</code></pre>
<p>For a small change in <code>theta</code>, <code>delta_theta</code>, the new <code>argmin</code> of <code>f</code> occurs when<br>
<code>d_f (delta_x) / d_delta_x</code> is zero.  Setting this derivative of the above expression<br>
to zero and solving for <code>delta_x</code> we get:</p>
<pre><code class="lang-plaintext">delta_x = -delta_theta * d^2_f/d_x d_theta / d^2_f/d_x^2
</code></pre>
<p>This <code>delta_x</code> is linear in <code>delta_theta</code>, so its gradient with respect to <code>delta_theta</code><br>
is simply <code>-d^2_f/d_x d_theta / d^2_f/d_x^2</code> and this is the desired gradient of<br>
<code>x_star = argmin_x f (x; theta)</code> with respect to <code>theta</code>.</p>
<p>Here is a pytorch script that performs this computation for a specific concrete<br>
network:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2023)

f = torch.nn.Sequential (   # network f, with fixed randomly initialized parameters (theta)
    torch.nn.Linear (1, 3),
    torch.nn.Tanh(),
    torch.nn.Linear (3, 1, bias = False)
)

x = torch.tensor ([0.0], requires_grad = True)   # vary x to minimize f(x)

# use pytorch optimizer to find argmin of f(x)
opt = torch.optim.Adam ([x], lr = 0.5)
for  i in range (350):
    l = f (x)
    opt.zero_grad()
    l.backward()
    opt.step()

# x is now x_star, the argmin of f(x) (for fixed theta)
x_star = x.detach().clone()   # save a copy for later
print ('x_star:', x_star.item())
print ('l:', l.item())

# compute second partials, d^2_f / d_x^2 an d^2_f / d_x d_theta, evaluated at x = x_star

x.grad = None
l = f (x)
d_x = torch.autograd.grad (l, x, retain_graph=True, create_graph=True)[0]   # first derivative of f with respect to x
print ('d_x:', d_x)                                                         # d_x is zero because x_star is a stationary point
d2_x = torch.autograd.grad (d_x, x, retain_graph = True)[0]                 # second derivative of f with respect to x
print ('d2_x:', d2_x)                                                       # d2_x is positive because that stationary point is a minimum

d_x_d_th = torch.autograd.grad (d_x, f.parameters())                        # mixed second partials
print ('d_x_d_th:')
for  t in d_x_d_th:
    print (t)

d_x_star_d_weight_1 = -d_x_d_th[0] / d2_x                                   # gradient with respect to weight of first Linear
d_x_star_d_bias_1 = -d_x_d_th[1] / d2_x                                     # gradient with respect to bias of first Linear
d_x_star_d_weight_2 = -d_x_d_th[2] / d2_x                                   # gradient with respect to weight of second Linear

print ('d_x_star_d_theta:')
print (d_x_star_d_weight_1)
print (d_x_star_d_bias_1)
print (d_x_star_d_weight_2)

# numerically check one value of gradient
delta = 1.e-4
with torch.no_grad():
    f[0].weight[1, 0] += delta

# recompute x_star for perturbed theta
for  i in range (125):
    l = f (x)
    opt.zero_grad()
    l.backward()
    opt.step()

grad = d_x_star_d_weight_1[1, 0]
grad_num = (x - x_star) / delta

print ('grad:', grad.item())
print ('grad_num:', grad_num.item())
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">1.13.1
x_star: 4.385826110839844
l: -0.6181716918945312
d_x: tensor([-3.3542e-09], grad_fn=&lt;SqueezeBackward1&gt;)
d2_x: tensor([0.0069])
d_x_d_th:
tensor([[-0.0194],
        [ 0.0611],
        [ 0.0005]])
tensor([0.0115, 0.0190, 0.0001])
tensor([[-0.0941,  0.0303,  0.0002]])
d_x_star_d_theta:
tensor([[ 2.8206],
        [-8.8939],
        [-0.0791]])
tensor([-1.6800, -2.7720, -0.0207])
tensor([[13.7074, -4.4159, -0.0283]])
grad: -8.89388656616211
grad_num: -8.88824462890625
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

medvisioner: <p>Hi Frank, many thanks for your wonderful answer!</p>
<p>I notice you jitter the \theta to \theta+\delta, in order to estimate the gradient w.r.t. \theta.</p>
<p>What is the difference between your method and the finite difference technique?  Why do you need to expand f to the second derivative at first?</p>
<p>The finite difference technique directly jitters \theta to \theta+\delta and minimizes the f_(\theta+\delta) to obtain x*_new, then estimates the gradient via x*_old - x*_new / \delta.</p>

KFrank: <p>Hi Med!</p>
<aside class="quote no-group" data-username="medvisioner" data-post="3" data-topic="173060" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/45deac/48.png" class="avatar"> medvisioner:</div>
<blockquote>
<p>I notice you jitter the \theta to \theta+\delta, in order to estimate the gradient w.r.t. \theta.</p>
</blockquote>
</aside>
<p>Just to be clear, I did two separate things.  The main computation used the<br>
second-order expansion of <code>f (x;, theta)</code>, plus autograd’s ability to compute<br>
second derivatives, to compute the gradient of <code>x_star</code> with respect to <code>theta</code>.</p>
<p>There are no finite differences here – neither in the expression for that gradient<br>
nor in autograd’s computation of the second partials.</p>
<p>Then, separately, I performed a numerical finite-difference computation (“jitter<br>
the \theta to \theta+\delta”) of just one element of the full gradient, that is, I<br>
estimated the derivative of <code>x_star</code> with respect to just one element of one<br>
of the parameter tensors.  This was just a numerical cross-check to show that<br>
the first computation was correct.</p>
<aside class="quote no-group">
<blockquote>
<p>What is the difference between your method and the finite difference technique?</p>
</blockquote>
</aside>
<p>As noted above, the main computation has nothing to do with finite differences,<br>
while the second computation – the numerical cross-check – is indeed a standard<br>
finite-difference estimate of the derivative.</p>
<aside class="quote no-group">
<blockquote>
<p>Why do you need to expand f to the second derivative at first?</p>
</blockquote>
</aside>
<p><code>x_star</code> – being the argmin – sits at the bottom of the “bowl” defined by <code>f (x)</code>.<br>
The second derivative of <code>f (x)</code> with respect to <code>x</code> tells us the curvature of the<br>
bottom of the bowl – that is, is the bowl relatively flat or more steep?</p>
<p>The second-order mixed partials of <code>f (x; theta)</code> with respect to <code>x</code> and <code>theta</code><br>
give us information about how changes in <code>theta</code> change the bowl, and, in<br>
particular, how they change the location of the bottom of the bowl (which is to<br>
say, how they change <code>x_star</code>).</p>
<p>I think about it like this:  The mixed partials tell us how hard changes in <code>theta</code><br>
“push” on <code>x_star</code>, while the second derivative with respect to <code>x</code> tells us how<br>
firmly <code>x_star</code> resists that push – a steeper bowl means that <code>x_star</code> won’t<br>
move as far in response to that push, while a flatter bowl means that <code>x_star</code><br>
will move farther.  Combined together, the two tell us how far <code>x_star</code> moves<br>
when <code>theta</code> changes, or, in more mathematical language, what the gradient<br>
of <code>x_star</code> with respect to <code>theta</code> is.</p>
<p>Best.</p>
<p>K. Frank</p>

medvisioner: <p>Thanks for the explanation. Problem solved!</p>

XuehaiPan: <blockquote>
<pre data-code-wrap="python"><code class="lang-python">d_x_star_d_weight_1 = -d_x_d_th[0] / d2_x                                   # gradient with respect to weight of first Linear
d_x_star_d_bias_1 = -d_x_d_th[1] / d2_x                                     # gradient with respect to bias of first Linear
d_x_star_d_weight_2 = -d_x_d_th[2] / d2_x                                   # gradient with respect to weight of second Linear
</code></pre>
</blockquote>
<p>When you involve with multi-dimensional inputs of <code>x</code>, the division <code>... / d2_x</code> will not work. You should use the inverse of a Jacobian matrix here.</p>
<p>FYI, you could try out <a href="https://github.com/metaopt/torchopt" rel="noopener nofollow ugc"><code>torchopt</code></a>, see the <a href="https://torchopt.readthedocs.io/en/latest/implicit_diff/implicit_diff.html" rel="noopener nofollow ugc">Implicit Gradient Differentiation</a> documentation for more details.</p>

medvisioner: <p>Hi, Frank.</p>
<p>I tried your code and was able to reproduce the result for random seed 2023.</p>
<p>However, when I set the seed to 456, the analytical gradient and numerical one is quite different. <img src="https://discuss.pytorch.org/images/emoji/apple/smiling_face_with_tear.png?v=12" title=":smiling_face_with_tear:" class="emoji" alt=":smiling_face_with_tear:" loading="lazy" width="20" height="20"></p>
<p>Any idea why this could happen?</p>
<p>My code:</p>
<pre><code class="lang-auto">import torch
print (torch.__version__)

_ = torch.manual_seed (456)

f = torch.nn.Sequential (   # network f, with fixed randomly initialized parameters (theta)
    torch.nn.Linear (1, 3),
    torch.nn.Tanh(),
    torch.nn.Linear (3, 1, bias = False)
)

x = torch.tensor ([0.0], requires_grad = True)   # vary x to minimize f(x)

# use pytorch optimizer to find argmin of f(x)
opt = torch.optim.Adam ([x], lr = 0.5)
for  i in range (350):
    l = f (x)
    opt.zero_grad()
    l.backward()
    opt.step()

# x is now x_star, the argmin of f(x) (for fixed theta)
x_star = x.detach().clone()   # save a copy for later
print ('x_star:', x_star.item())
print ('l:', l.item())

# compute second partials, d^2_f / d_x^2 an d^2_f / d_x d_theta, evaluated at x = x_star

x.grad = None
l = f (x)
d_x = torch.autograd.grad (l, x, retain_graph=True, create_graph=True)[0]   # first derivative of f with respect to x
print ('d_x:', d_x)                                                         # d_x is zero because x_star is a stationary point
d2_x = torch.autograd.grad (d_x, x, retain_graph = True)[0]                 # second derivative of f with respect to x
print ('d2_x:', d2_x)                                                       # d2_x is positive because that stationary point is a minimum

d_x_d_th = torch.autograd.grad (d_x, f.parameters())                        # mixed second partials
print ('d_x_d_th:')
for  t in d_x_d_th:
    print (t)

d_x_star_d_weight_1 = -d_x_d_th[0] / d2_x                                   # gradient with respect to weight of first Linear
d_x_star_d_bias_1 = -d_x_d_th[1] / d2_x                                     # gradient with respect to bias of first Linear
d_x_star_d_weight_2 = -d_x_d_th[2] / d2_x                                   # gradient with respect to weight of second Linear

print ('d_x_star_d_theta:')
print (d_x_star_d_weight_1)
print (d_x_star_d_bias_1)
print (d_x_star_d_weight_2)

# numerically check one value of gradient
delta = 1.e-4
with torch.no_grad():
    f[0].weight[1, 0] += delta

# recompute x_star for perturbed theta
for  i in range (125):
    l = f (x)
    opt.zero_grad()
    l.backward()
    opt.step()

grad = d_x_star_d_weight_1[1, 0]
grad_num = (x - x_star) / delta

print ('grad:', grad.item())
print ('grad_num:', grad_num.item())
</code></pre>
<p>My output:</p>
<pre><code class="lang-auto">1.10.0
x_star: -12.196703910827637
l: -0.7395662069320679
d_x: tensor([0.0004], grad_fn=&lt;SqueezeBackward1&gt;)
d2_x: tensor([0.0003])
d_x_d_th:
tensor([[-0.0081],
        [ 0.0000],
        [ 0.0002]])
tensor([ 7.4753e-04, -0.0000e+00, -2.0273e-05])
tensor([[6.7220e-04, 0.0000e+00, 2.5887e-05]])
d_x_star_d_theta:
tensor([[30.9327],
        [-0.0000],
        [-0.8667]])
tensor([-2.8593,  0.0000,  0.0775])
tensor([[-2.5712, -0.0000, -0.0990]])
grad: -0.0
grad_num: -6002.197265625
</code></pre>

medvisioner: <p>Hello, Frank and Xuehai:</p>
<p>I implement a toy model to estimate dx_dtheta where x is a high-dimensional vector.</p>
<p>Specifically, I use dx_dtheta = - H_f[x]^-1 * d2f_fxdtheta, where H_f[x]^-1 is the inverse Hessian and the d2f_fxdtheta is the mixed partial.</p>
<p>However, the dx_dtheta I estimated is very stange. They are about 10^15 large.<br>
As a contrast, the result estimated via finite difference is about 1.</p>
<p>My codes and ouputs are attached below. Any idea which part is wrong?</p>
<p>Codes:</p>
<pre><code class="lang-auto">import torch
_ = torch.manual_seed(123)
torch.set_default_dtype(torch.float64)

print(torch.__version__)

f = torch.nn.Sequential (   # network f, with fixed randomly initialized parameters (theta)
    torch.nn.Linear (3, 2),
    torch.nn.Tanh(),
    torch.nn.Linear (2, 1, bias = False)
)

x = torch.tensor ([0.0, 0.0, 0.0], requires_grad = True)   # vary x to minimize f(x)

# use pytorch optimizer to find argmin of f(x)
opt = torch.optim.Adam ([x], lr = 0.5)
for  i in range (350):
    l = f (x)
    opt.zero_grad()
    l.backward()
    opt.step()
    

# x is now x_star, the argmin of f(x) (for fixed theta)
x_star = x.detach().clone()   # save a copy for later
print ('x_star:', x_star)
print ('l:', l.item())


# compute second partials, d^2_f / d_x^2 an d^2_f / d_x d_theta, evaluated at x = x_star
x.grad = None
l = f (x)
d_x = torch.autograd.grad (l, x, retain_graph=True, create_graph=True)[0]   # first derivative of f with respect to x
print ('d_x:', d_x)   
d2_x = torch.stack([torch.autograd.grad (d_x[ind], x, retain_graph=True, create_graph=True)[0] for ind in range(3)])
print ('d2_x:', d2_x)   


# only look at the first fc layer's weight
theta = list(f.parameters())[0]

d_f_d_x_d_th = torch.stack([torch.autograd.grad(d_x[ind], theta, retain_graph=True, create_graph=True)[0] for ind in range(3)]).reshape(*x.shape,*theta.shape)                        # mixed second partials
print ('d_f_d_x_d_th:', d_f_d_x_d_th)

# use the IFT here
d_x_d_th = - torch.matmul(torch.linalg.inv(d2_x), d_f_d_x_d_th.reshape(3,-1))
print('d_x_d_th:', d_x_d_th)

# numerically check one value of gradient
# set delta to 1e-6, 1e-8, 1e-10 gives the same numerical values, so it is quite stable

x = torch. Tensor ([0.0,0.0,0.0], requires_grad = True)   # vary x to minimize f(x)
opt = torch.optim.Adam ([x], lr = 0.5)

delta = 1.e-10
with torch.no_grad():
    f[0].weight[0, 0] += delta

# recompute x_star for perturbed theta
for  i in range (350):
    l = f (x)
    opt.zero_grad()
    l.backward()
    opt.step()

grad_num = (x - x_star) / delta

print('analycal: ',d_x_d_th[:,0])
print('numerial: ',grad_num)
</code></pre>
<p>Outputs:</p>
<pre><code class="lang-auto">1.10.0
x_star: tensor([-10.2784,   9.1450,  -8.9501])
l: -1.2608992210123384
d_x: tensor([ 1.3880e-04, -1.9742e-04,  2.9152e-05], grad_fn=&lt;SqueezeBackward1&gt;)
d2_x: tensor([[ 2.0108e-04,  7.6194e-05, -2.1971e-05],
        [ 7.6194e-05,  2.3750e-04, -4.5307e-05],
        [-2.1971e-05, -4.5307e-05,  8.9560e-06]], grad_fn=&lt;StackBackward0&gt;)
d_f_d_x_d_th: tensor([[[ 0.0015, -0.0010,  0.0010],
         [ 0.0036, -0.0036,  0.0035]],

        [[ 0.0043, -0.0035,  0.0038],
         [ 0.0003, -0.0007,  0.0002]],

        [[-0.0008,  0.0007, -0.0003],
         [-0.0002,  0.0002, -0.0006]]], grad_fn=&lt;ReshapeAliasBackward0&gt;)
d_x_d_th: tensor([[ 2.0818e+15,  8.7662e+15,  4.9454e+16, -2.3087e+15, 
-9.7215e+15,
         -5.4843e+16],
        [ 8.7662e+15,  3.6913e+16,  2.0824e+17, -9.7215e+15, -4.0935e+16,
         -2.3093e+17],
        [ 4.9454e+16,  2.0824e+17,  1.1748e+18, -5.4843e+16, -2.3093e+17,
         -1.3028e+18]], grad_fn=&lt;NegBackward0&gt;)
analycal:  tensor([2.0818e+15, 8.7662e+15, 4.9454e+16], 
grad_fn=&lt;SelectBackward0&gt;)
numerial:  tensor([  0.9589, -12.7340,  15.6670], grad_fn=&lt;DivBackward0&gt;)
</code></pre>

XuehaiPan: <p>A reproduction script using <a href="https://github.com/metaopt/torchopt" rel="noopener nofollow ugc"><code>torchopt</code></a>’s implicit gradient:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn

import torchopt

class ImplicitModel(torchopt.nn.ImplicitMetaGradientModule):
    def __init__(self, mlp, x0):
        super().__init__()
        self.mlp = mlp
        self.x = nn.Parameter(x0.clone().detach_(), requires_grad=True)

    def objective(self):
        return self.mlp(self.x).squeeze()

    @torch.enable_grad()
    def solve(self):
        optimizer = torch.optim.Adam([self.x], lr=0.01)
        for i in range(1000):
            optimizer.zero_grad()
            loss = self.objective()
            loss.backward(inputs=[self.x])
            optimizer.step()

torch.manual_seed(456)

mlp = torch.nn.Sequential(
    torch.nn.Linear(1, 3),
    torch.nn.Tanh(),
    torch.nn.Linear(3, 1, bias=False),
)
x0 = torch.tensor([0.0])

model = ImplicitModel(mlp, x0)
model.solve()

model.x[0].backward()

for name, meta_param in model.named_meta_parameters():
    print(f'{name}.grad: {meta_param.grad}')
</code></pre>
<p>Output for seed 456:</p>
<pre><code class="lang-python">mlp.0.weight.grad: tensor([[13.5392],
        [-0.0203],
        [-3.2621]])
mlp.0.bias.grad: tensor([-3.8148,  0.0044,  0.7939])
mlp.2.weight.grad: tensor([[-3.9982e+00,  3.7859e-03, -1.0407e+00]])
</code></pre>
<p>Output for seed 2023:</p>
<pre><code class="lang-python">mlp.0.weight.grad: tensor([[ 2.6780],
        [-5.9816],
        [-0.1180]])
mlp.0.bias.grad: tensor([-0.8137, -2.4569, -0.0382])
mlp.2.weight.grad: tensor([[ 7.5996, -4.0433, -0.0522]])
</code></pre>

XuehaiPan: <p>Reproduction script for multi-dimension inputs:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn

import torchopt

class ImplicitModel(torchopt.nn.ImplicitMetaGradientModule):
    def __init__(self, mlp, x0):
        super().__init__()
        self.mlp = mlp
        self.x = nn.Parameter(x0.clone().detach_(), requires_grad=True)

    def objective(self):
        return self.mlp(self.x).squeeze()

    @torch.enable_grad()
    def solve(self):
        optimizer = torch.optim.Adam([self.x], lr=0.01)
        for i in range(1000):
            optimizer.zero_grad()
            loss = self.objective()
            loss.backward(inputs=[self.x])
            optimizer.step()

torch.manual_seed(123)

mlp = torch.nn.Sequential(
    torch.nn.Linear(3, 2),
    torch.nn.Tanh(),
    torch.nn.Linear(2, 1, bias=False),
)
x0 = torch.tensor([0.0, 0.0, 0.0])

for i in range(3):
    for param in mlp.parameters():
        param.detach_().requires_grad_()

    model = ImplicitModel(mlp, x0)
    model.solve()

    print(f'Derivative for x[{i}]')
    model.x[i].backward()
    for name, meta_param in model.named_meta_parameters():
        print(f'{name}.grad: {meta_param.grad}')
    print()
</code></pre>
<p>Output for seed 123:</p>
<pre><code class="lang-python">Derivative for x[0]
mlp.0.weight.grad: tensor([[  2.9149,  60.1467, -12.5878],
        [ -1.9555,   0.3308,  -1.7174]])
mlp.0.bias.grad: tensor([2.2414, 0.4061])
mlp.2.weight.grad: tensor([[2.6185, 0.6343]])

Derivative for x[1]
mlp.0.weight.grad: tensor([[ 47.1573, 334.5729, -46.8080],
        [-11.7130,   2.0863, -10.2923]])
mlp.0.bias.grad: tensor([4.7971, 2.4490])
mlp.2.weight.grad: tensor([[5.6041, 3.8254]])

Derivative for x[2]
mlp.0.weight.grad: tensor([[ 35.8865, 314.7516, -50.3962],
        [-10.7902,   1.8953,  -9.4800]])
mlp.0.bias.grad: tensor([6.6147, 2.2518])
mlp.2.weight.grad: tensor([[7.7275, 3.5174]])
</code></pre>

medvisioner: <p>Under the same setting, the finite difference estimates dx_dmlp.0.weight[0] to be <code>[0.1900, -4.9266,  5.9873]</code>, while torchopt gives <code>[-1.5724,-9.2159,-7.7948]</code>.</p>
<p>So basically, the inverse Hessian, torchopt, and finite difference gives totally different results? <img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=12" title=":joy:" class="emoji" alt=":joy:" loading="lazy" width="20" height="20"></p>
<p>I check, when <code>x=torch. Tensor([0.0])</code> is a scaler, both the inverse Hessian, torchopt, and finite difference gives the same gradient estimation.</p>

KFrank: <p>Hi Med!</p>
<aside class="quote no-group quote-modified" data-username="medvisioner" data-post="7" data-topic="173060" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/45deac/48.png" class="avatar"> medvisioner:</div>
<blockquote>
<p>However, when I set the seed to 456, the analytical gradient and numerical one is quite different. <img src="https://discuss.pytorch.org/images/emoji/apple/smiling_face_with_tear.png?v=12" title=":smiling_face_with_tear:" class="emoji" alt=":smiling_face_with_tear:" loading="lazy" width="20" height="20"></p>
<p>Any idea why this could happen?</p>
<p>…</p>
<p>My output:</p>
<pre><code class="lang-auto">x_star: -12.196703910827637
l: -0.7395662069320679
d_x: tensor([0.0004], grad_fn=&lt;SqueezeBackward1&gt;)
d2_x: tensor([0.0003])
</code></pre>
</blockquote>
</aside>
<p>First note that you haven’t converged to a minimum.  Your first derivative,<br>
although smallish, is not zero.  (Also your second derivative is relatively<br>
small, so if you were at a minimum, it would be pretty flat.)  The formula<br>
I used for the gradient of argmin relies on <code>x_star</code> actually being a minimum<br>
with <code>d_f/d_x</code> equaling zero when evaluated at <code>x_star</code>.</p>
<p>Not having converged to a minimum is not simply a consequence of not<br>
having run enough optimization steps.</p>
<p>Here are the details of what is going on:</p>
<p>The particular “model” I used in my example is not guaranteed to have a<br>
minimum.  When the <code>weight</code>s of the <code>Linear</code>s are initialized starting with a<br>
random seed of <code>2023</code>, the model does turn out to have a minimum.  But<br>
when initialized with a random seed of <code>456</code>, it does not.  As you optimize,<br>
<code>x</code> runs off to <code>-inf</code> and the <code>Tanh</code> “saturates” and becomes flat.</p>
<p>To make this clear, we can add a “stabilizing” term to the model that will<br>
ensure that it has a minimum.  So we run the same analysis on the <code>456</code><br>
model plus an <code>x**2</code> term and see that the finite-difference gradient and<br>
“analytic” gradient agree (as well as can be expected).</p>
<p>For completeness, here is the modified script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (456)   # different seed causes different random initialization of model

class FStab (torch.nn.Module):
    def __init__ (self):
        super().__init__()
        self.fseq = torch.nn.Sequential (   # network f, with fixed randomly initialized parameters (theta)
            torch.nn.Linear (1, 3),
            torch.nn.Tanh(),
            torch.nn.Linear (3, 1, bias = False)
        )
    def forward (self, x):                  # add stabilization term in forward
        return  self.fseq (x) + x**2 / 16

f = FStab()

x = torch.tensor ([0.0], requires_grad = True)   # vary x to minimize f(x)

# use pytorch optimizer to find argmin of f(x)
opt = torch.optim.Adam ([x], lr = 0.5)
for  i in range (350):
    l = f (x)
    opt.zero_grad()
    l.backward()
    opt.step()

# x is now x_star, the argmin of f(x) (for fixed theta)
x_star = x.detach().clone()   # save a copy for later
print ('x_star:', x_star.item())
print ('l:', l.item())

# compute second partials, d^2_f / d_x^2 an d^2_f / d_x d_theta, evaluated at x = x_star

x.grad = None
l = f (x)
d_x = torch.autograd.grad (l, x, retain_graph=True, create_graph=True)[0]   # first derivative of f with respect to x
print ('d_x:', d_x)                                                         # d_x is zero because x_star is a stationary point
d2_x = torch.autograd.grad (d_x, x, retain_graph = True)[0]                 # second derivative of f with respect to x
print ('d2_x:', d2_x)                                                       # d2_x is positive because that stationary point is a minimum

d_x_d_th = torch.autograd.grad (d_x, f.parameters())                    # mixed second partials
print ('d_x_d_th:')
for  t in d_x_d_th:
    print (t)

d_x_star_d_weight_1 = -d_x_d_th[0] / d2_x                                   # gradient with respect to weight of first Linear
d_x_star_d_bias_1 = -d_x_d_th[1] / d2_x                                     # gradient with respect to bias of first Linear
d_x_star_d_weight_2 = -d_x_d_th[2] / d2_x                                   # gradient with respect to weight of second Linear

print ('d_x_star_d_theta:')
print (d_x_star_d_weight_1)
print (d_x_star_d_bias_1)
print (d_x_star_d_weight_2)

# numerically check one value of gradient
delta = 1.e-4
with torch.no_grad():
    f.fseq[0].weight[1, 0] += delta

# recompute x_star for perturbed theta
for  i in range (125):
    l = f (x)
    opt.zero_grad()
    l.backward()
    opt.step()

grad = d_x_star_d_weight_1[1, 0]
grad_num = (x - x_star) / delta

print ('grad:', grad.item())
print ('grad_num:', grad_num.item())
</code></pre>
<p>And its output:</p>
<pre><code class="lang-plaintext">1.13.1
x_star: -1.146886944770813
l: -0.24911892414093018
d_x: tensor([0.], grad_fn=&lt;AddBackward0&gt;)
d2_x: tensor([0.2937])
d_x_d_th:
tensor([[ 0.6172],
        [ 0.1321],
        [-0.2920]])
tensor([-0.0671, -0.2309, -0.0733])
tensor([[ 0.3524, -0.2288,  0.4708]])
d_x_star_d_theta:
tensor([[-2.1014],
        [-0.4498],
        [ 0.9942]])
tensor([0.2285, 0.7860, 0.2496])
tensor([[-1.1997,  0.7790, -1.6029]])
grad: -0.4498182535171509
grad_num: -0.4506111145019531
</code></pre>
<p>Note that when evaluated at the (successfully-converged) <code>x_star</code>, <code>d_x</code> is<br>
zero and <code>d2_x</code> is comfortably positive.</p>
<p>Best.</p>
<p>K. Frank</p>

medvisioner: <p>Thanks Frank. You explanation is very clear.</p>