How can I use torch's `VHP` routine inside of the `step` method when defining a custom optimizer?

username: aryamanjeendgar
cooked: <p>So I have been playing around with implementing an optimizer that leverages the hessian-vector product in it’s update rule inside of torch — I am keen on using the existing <a href="https://pytorch.org/docs/stable/generated/torch.autograd.functional.vhp.html" rel="noopener nofollow ugc"><code>VHP</code></a> for the same, but it’s really awkward API makes it impossible to be used inside of the <code>step</code> method when defining a custom optimizer class in torch.</p>
<p>Here are the problems that I currently have with the existing <code>VHP</code>interface:</p>
<ol>
<li>One typically does not pass in any argument to the <code>step</code>method when making the call to <code>optimizer.step()</code>, <em>except</em> in the case of algorithms like <code>L-BFGS</code> which require a <code>closure</code>argument since they re-evaluate the loss function multiple times every iteration. Given the API for the VHP, a typical call to the method would look something like: <code>VHP(loss_func, (model_output_tensor, gt_labels/values_tensor), v)</code> (I have gripes about this signature as well, which I explain in the next point). But <code>step</code> does <em>not</em> typically have access to <em>any</em> of these arguments! <code>step</code> doesn’t have any idea about <code>loss_func, model_output_tensor, gt_labels</code> — so just to get it to work I decided to dump all of these arguments inside of <code>closure</code> (horrible choice, but that seemed to be the only way out) but then…</li>
<li>It seems like the way <code>VHP</code> treats the <code>inputs</code> tuple is as separate inputs to the function that are to be evaluated independently (that is why they explicitly impose <code>len(inputs) == len(v)</code>), but loss functions almost always take in vector inputs and I have no clue how one could pass in <em>two</em> vector inputs in a single tensor (except maybe by vertically stacking the inputs, but I am unsure if <code>VHP</code> would parse the inputs that way).</li>
</ol>
<p>TLDR: The API for torch’s <code>VHP</code> method is <em>really</em> hard to work with when defining a custom optimizer in torch — are there any interesting alternatives for computing the Hessian-Vector-Product that do <em>not</em> run into the pitfalls that I mentioned above? As far as I am aware <code>JAX</code> also has an identical API for their own <code>HVP</code> method.</p>
<p>Thank you for your time!</p>

username: soulitzer
cooked: <p>(1) looks more like a complaint to the optimizer interface haha, but using a closure unfortunately is what you have to do yes<br>
(2) I’m not sure I understand. Isn’t hessian inputs x inputs, and so you need v to also be inputs in order to compute vhp? so that is why len(inputs) == len(v) is necessary.</p>

username: aryamanjeendgar
cooked: <ol>
<li>Haha, I agree XD.</li>
<li>For this point, consider the example of the following function:</li>
</ol>
<pre><code class="lang-auto">def test_func_multi(x: torch.Tensor, y:torch.Tensor):
    return (2 * y ** 3 - 3 * x ** 2).sum()
</code></pre>
<p>I do not understand how to use the existing VHP API for this (and it doesn’t help that the documentation doesn’t have such an example) — like, I would want to pass in vector inputs for both <code>x</code> and <code>y</code> in this case, but if I were to pass in the same as a tuple as: <code>VHP(test_func_multi, (x_in, y_in), v)</code> where the semantics I hope for are that <code>(x_in, y_in)</code> are used for instantiating <code>test_func_multi</code>, but this throws an error about the lengths of the inputs and the vector not being equal. I am interested in a function like this because that is the signature of <em>every</em> loss function.</p>
<p>One way around this could be to partially instantiate the <code>test_func_multi</code> function using <code>functools.partial</code> (since for a loss function, the <code>labels</code> are anyways tensors that have <code>requires_grad=False</code>)</p>

username: aryamanjeendgar
cooked: <p>I had a separate question as well, I would be really thankful if you could help :).</p>
<p>Since the posting of this question, I have figured out a way to compute the <code>HVP</code> for my use-case inside of the <code>step</code> method <em>without</em> having to resort to passing in additional arguments inside of <code>closure</code> (through some <code>autograd</code> magic XD). The optimizer seems to run fine for synthetic experiments (like on static functions e.g. rosenbrock etc.), but it seems to throw a really opaque error after running for a few iterations when running for simple NN’s (in this case a simple CNN over CIFAR10).</p>
<p>This is the error that I am getting:</p>
<pre><code class="lang-auto">E               RuntimeError: A view was created in no_grad mode and is being modified inplace with grad mode enabled. Given that this use case is ambiguous and error-prone, it is forbidden. You can clarify your code by moving both the view and the inplace either both inside the no_grad block (if you don't want the inplace to be tracked) or both outside (if you want the inplace to be tracked).
</code></pre>
<p>I found <a href="https://github.com/pytorch/pytorch/issues/99968" rel="noopener nofollow ugc">this relevant torch issue</a> (where I believe you have also replied to the conversation?)</p>
<p>I get this error on a line of the <code>step</code> method where I am making updates to a <code>state</code> variable of the optimizer (this particular <code>state</code> variable also happens to be involved in the HVP) — and interestingly, I get this error <em>after</em> the algorithm has already run for a few iterations, making this <em>really</em> hard to debug. The update is very simple as well, it’s just <code>T[T &gt; constant] = constant</code> (where <code>T</code> is the state variable in question).</p>
<p>Would you have any comments for what I could do in this situation?</p>
<p>Thank you so much for your help!</p>

username: soulitzer
cooked: <p>This might be a dumb question, but why won’t you have two v’s in your example?</p>

username: soulitzer
cooked: <p>Is there any reason the inplace update to state needs to be done in grad mode? Doing the update to state in no-grad would be one way of resolving this.</p>