Parameter update to minimize a loss that doesn't directly depend on said parameters

peustr: <p>This sounds counterintuitive but let me explain.</p>
<p>I have 2 nn.Modules, let’s call them <code>model</code> and <code>input_generator</code>.<br>
<code>input_generator</code> accepts a data point <code>x</code> as input, and outputs a vector <code>x_new</code>. <code>x_new</code> is in turn passed to the <code>model</code>, and <code>model</code> is trained to minimize the cross entropy loss between <code>model(x_new)</code> and <code>target</code>. In code:</p>
<pre><code class="lang-python"># assume: input variable 'x', class 'target', optimizer on model.parameters
x_new = input_generator(x)
model_optimizer.zero_grad()
logits = model(x_new)
loss = F.cross_entropy(logits, target)
model_optimizer.step()
</code></pre>
<p>So far so good. Now the question is, how is <code>input_generator</code> trained? Well, I want <code>input_generator</code> to generate inputs that minimize a <strong>different</strong> objective for <code>model</code>. Say we have a completely unrelated input <code>x_other</code>. I want <code>input_generator</code> to generate <code>x_new</code> inputs to minimize the cross-entropy loss between <code>model(x_other)</code> and <code>target</code>. Expanding on the previous code snippet:</p>
<pre><code class="lang-python">x_new = input_generator(x)
model_optimizer.zero_grad()
logits = model(x_new)
loss = F.cross_entropy(logits, target)
loss.backward()
model_optimizer.step()
# assume: unrelated input 'x_other', class 'other_target', optimizer on input_generator.parameters
ig_optimizer.zero_grad()
other_logits = model(x_other)
other_loss = F.cross_entropy(other_logits , other_target)
other_loss.backward()
ig_optimizer.step()  # this won't work of course
</code></pre>
<p>Now, you see the problem. There is no way to optimize <code>input_generator</code> with my “off-policy” objective, because <code>input_generator</code> does not play absolutely any role in the creation of <code>x_other</code>.</p>
<p>So, my dear community, is there a way to say: “Please, <code>input_generator</code>, generate <code>x_new</code> in a way that minimizes the cross-entropy loss between <code>model</code> and a seemingly unrelated input, once <code>model</code> has taken a gradient step with your generated <code>x_new</code>.”?</p>
<p>Thank you and all the best.</p>

srishti-git1110: <p>Hi <a class="mention" href="/u/peustr">@peustr</a><br>
I don’t think there is a way to do that, that is at least theoretically right.</p>
<p><code>other_loss</code> is a function of <code>other_logits</code> (and <code>other_target</code> but that is a constant hence not optimizable).<br>
<code>other_logits</code> in turn is a function of <code>model</code> parameters only. So, when you call <code>other_loss.backward()</code>, gradients of <code>other_loss</code> with respect to <code>model</code> parameters shall be computed.</p>
<p><code>ig_optimizer.step()</code> needs the gradients of <code>other_loss</code> wrt to the parameters of <code>input_generator</code> to be able to update them.<br>
Parameters of <code>input_generator</code> have nothing to do with what value is taken by <code>other_loss</code>. That to say, you really cannot “optimise” the parameters of <code>input_generator</code> for <code>other_loss</code> to take an optimised value.</p>
<p>Let’s think of this using an example-</p>
<pre><code class="lang-auto">z = y*y
x = torch.tensor([1.0], requires_grad=True)
</code></pre>
<p>z is a function of tensor y. To minimize z, x is simply irrelevant.</p>

peustr: <p>Hi <a class="mention" href="/u/srishti-git1110">@srishti-git1110</a>, thank you for your reply.</p>
<p>I understand completely what you are saying, but the example you provide is not a direct analogy.</p>
<p>If we think about it, there exists a path from <code>x</code> to <code>other_loss</code> that passes through <code>x_new</code>. The path is as follows:</p>
<ul>
<li>
<code>x_new</code> is generated by <code>input_generator(x)</code> e.g., a function with parameters φ</li>
<li>
<code>x_new</code> is given to <code>model</code>, produces <code>loss</code>
</li>
<li>
<code>loss</code> is used to update <code>model</code> parameters θ_t to θ_{t+1}</li>
<li>given a constant input <code>x_other</code>, <code>model</code> with parameters θ_{t+1} produces <code>other_loss</code>
</li>
<li>what is the <code>x_new</code> that would minimize <code>other_loss</code>, if <code>model</code> with parameters θ_t had used it to perform the gradient update instead?</li>
<li>update φ so that it produces a better <code>x_new</code> in the next iteration</li>
</ul>
<p>I am aware that this may not be possible (although it seems like it could be, with second-order gradients or something); I am just demonstrating that the example you provided does not represent my problem fully.</p>
<p>Thanks again!</p>

srishti-git1110: <p><a class="mention" href="/u/peustr">@peustr</a><br>
Yes, I included that example just to convey that we cannot traditionally optimise a function wrt some variable that it mathematically does not depend on (directly or indirectly).</p>
<p>But, what you detailed in the bullet points makes sense to me. Thanks for the clarification.</p>
<p>The task now will be to put these points in form of mathematical equations to get a sense of how autograd could be utilised to update input_generator’s parameters.<br>
I’ll get back to you on this if I’m able to figure something out.</p>

KFrank: <p>Hi Panagiotis!</p>
<aside class="quote no-group" data-username="peustr" data-post="1" data-topic="161099" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/peustr/48/22092_2.png" class="avatar"> peustr:</div>
<blockquote>
<p>So, my dear community, is there a way to say: “Please, <code>input_generator</code>, generate <code>x_new</code> in a way that minimizes the cross-entropy loss between <code>model</code> and a seemingly unrelated input, once <code>model</code> has taken a gradient step with your generated <code>x_new</code>.”?</p>
</blockquote>
</aside>
<p>If I understand your use case correctly, you would like to train <code>input_generator</code><br>
so that it produces samples of <code>x_new</code> which if used to train <code>model</code> with the<br>
loss function <code>cross_entropy (model (x_new), target)</code> then <code>model</code> will<br>
learn (from being trained with <code>x_new</code>) to minimize<br>
<code>cross_entropy (model (x_other), target)</code>.</p>
<p>Based on this understanding, one approach would be to train <code>input_generator</code><br>
to produce <code>x_new</code> such that <code>cross_entropy (model (x_new), target)</code> will<br>
have a gradient that trains <code>model</code> in the way you want.  To do this, you could<br>
use the gradient of <code>cross_entropy (model (x_other), target)</code> as the<br>
“ground-truth” target gradient for the gradient produced by<br>
<code>cross_entropy (model (x_new), target)</code> and then train <code>input_generator</code><br>
to produce samples of <code>x_new</code> that produce such gradients.</p>
<p>In the example below, <code>target_grad_weight</code> and <code>target_grad_bias</code> are the<br>
“target” gradients for the <code>weight</code> and <code>bias</code> of a simple <code>Linear</code> <code>model</code>.<br>
We then compute <code>x_new = input_generator (x)</code> (also for a simple <code>Linear</code><br>
<code>input_generator</code>) and compute <code>gen_grads</code>, the gradient with respect to<br>
<code>model</code> of <code>cross_entropy (model (x_new), target)</code>,  We then use the<br>
mean-squared-error between <code>gen_grads</code> and the target gradients as the<br>
loss function with which to train <code>input_generator</code>.</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; model = torch.nn.Linear (3, 5)
&gt;&gt;&gt; input_generator = torch.nn.Linear (2, 3)
&gt;&gt;&gt; x_other = torch.randn (1, 3)
&gt;&gt;&gt; target = torch.randint (5, (1,))
&gt;&gt;&gt;
&gt;&gt;&gt; loss_model_other = torch.nn.functional.cross_entropy (model (x_other), target)
&gt;&gt;&gt; loss_model_other.backward()
&gt;&gt;&gt; target_grad_weight = model.weight.grad.clone()
&gt;&gt;&gt; target_grad_bias = model.bias.grad.clone()
&gt;&gt;&gt; model.zero_grad()
&gt;&gt;&gt;
&gt;&gt;&gt; x = torch.randn (1, 2)
&gt;&gt;&gt; x_new = input_generator (x)
&gt;&gt;&gt; loss_model_new = torch.nn.functional.cross_entropy (model (x_new), target)
&gt;&gt;&gt; gen_grads = torch.autograd.grad (loss_model_new, (model.weight, model.bias), create_graph = True)
&gt;&gt;&gt;
&gt;&gt;&gt; loss_input_generator  = torch.nn.functional.mse_loss (gen_grads[0], target_grad_weight)
&gt;&gt;&gt; loss_input_generator += torch.nn.functional.mse_loss (gen_grads[1], target_grad_bias)
&gt;&gt;&gt;
&gt;&gt;&gt; loss_input_generator.backward()
&gt;&gt;&gt; input_generator.bias.grad
tensor([ 0.3475, -0.2202, -0.0620])
</code></pre>
<p>We do, in fact, compute a second-order derivative of a sort in that we<br>
compute the gradient of a function of a gradient and we use<br>
<code>autograd.grad (create_graph = True)</code> as an intermediate step in<br>
doing so.</p>
<p>Best.</p>
<p>K. Frank</p>

peustr: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a>, thank you so much for your reply, this is a very fresh perspective for me, I can’t wait to go to the office on Monday and try it out!<br>
I will mark your response as the solution for now, because it makes total sense to me that it should work, but I will make sure to leave an update here soon. Cheers!</p>