What is the purpose of the `closure` argument for `optimizer.step`?

username: Dominik
cooked: <p>Considering <a href="https://pytorch.org/docs/stable/optim.html" rel="noopener nofollow ugc">the docs</a> I realized that every <code>optimizer.step</code> method has an (optional) <code>closure</code> argument (for LBFGS it’s even required). The info text is:</p>
<blockquote>
<p>A closure that reevaluates the model and returns the loss.</p>
</blockquote>
<p>For me it is not completely clear which steps should be taken in this closure function. I reviewed some examples (for example <a href="https://pytorch.org/tutorials/advanced/neural_style_tutorial.html#gradient-descent" rel="noopener nofollow ugc">here</a> and <a href="https://github.com/leongatys/PytorchNeuralStyleTransfer/blob/master/NeuralStyleTransfer.ipynb" rel="noopener nofollow ugc">there</a>) and it looked like that within the closure we should:</p>
<ul>
<li>Zero the gradients</li>
<li>Compute the loss</li>
<li>Backprop on the loss</li>
<li>Return the loss</li>
</ul>
<p>So my questions are:</p>
<ol>
<li>Is the above list of (minimum required) steps that should be taken within the closure function correct and complete?</li>
<li>What is the purpose of returning the loss? What does the optimizer do with it? Should we return the <code>loss</code> that we backproped on or create a new one (the two examples handle this situation differently; maybe it doesn’t matter)?</li>
<li>For an optimizer for which the closure argument is optional, is <code>closure(); optimizer.step()</code> similar to <code>optimizer.step(closure)</code>?</li>
</ol>
<p>Thanks for your input!</p>

username: JTiC
cooked: <p>Have you known it? I have the similar question.</p>

username: Dominik
cooked: <p>I’m not completely sure about the internals but I’ve used <code>closure</code> that way successfully. So to comment on my own questions:</p>
<ol>
<li>I performed those steps in the <code>closure</code> and it worked.</li>
<li>Some optimizers (e.g. <a href="https://pytorch.org/docs/stable/optim.html#torch.optim.LBFGS" rel="nofollow noopener">LBFGS</a>) will terminate depending on the value of the <code>loss</code>, so that’s why it needs access to it and hence it should be returned.</li>
<li>I suppose those two versions are similar.</li>
</ol>
<p>Some optimizers (again LBFGS) need to evaluate the model multiple times hence the <code>closure</code> definition makes that possible.</p>
<p>Anyway I would be glad if someone with more in-depth knowledge could comment on these questions and clarify the situation.</p>

username: janeyx99
cooked: <p>This is definitely an ultra late reply, but I have just also come across this question myself. After looking in the code a bit, I am hoping to add a tad more clarity, but a researcher would know more!</p>
<p>The closure is mostly, if not entirely, only needed for LBFGS. I have not seen another use case for it with PyTorch where the typical way without the closure was not sufficient. (By the typical way, I mean the training loop of <code>zero_grad</code>, forward to get loss, backprop, <code>optimizer.step</code>.)</p>
<ol>
<li>
<p>In LBFGS, the closure is used to return the loss and update the gradients, so those two things would be the minimum requirement. In a typical model, the four steps you mentioned would suffice, but you could imagine cruder (though less useful) ways to update the grads and return a loss.</p>
</li>
<li>
<p>As you’ve concluded, LBFGS repeatedly <a href="https://github.com/pytorch/pytorch/blob/a2552d55215a0f5bc245ffcdb85024b92ff821ac/torch/optim/lbfgs.py#L280" rel="noopener nofollow ugc">recalculates the loss</a> and uses that and the updated grads to evaluate a direction until certain conditions are met. Internally, as of today, LBFGS immediately converts the losses to floats, so it doesn’t matter for LBFGS internals whether you return the literal loss tensor instance or a new one. Its <code>step</code> <em>will</em> return the original loss unprocessed, so if you intend to use the loss returned by <code>step</code>, then it would matter.</p>
</li>
<li>
<p>Yes, <code>closure(); optimizer.step()</code> is similar to <code>optimizer.step(closure)</code>, though it ultimately depends on how you define your closure. In most cases I’ve seen, this has been true, but I am not sure why people would use closures for anything other than LBFGS. One difference I did observe is that if you do pass in a closure, step will <a href="https://github.com/pytorch/pytorch/blob/a2552d55215a0f5bc245ffcdb85024b92ff821ac/torch/optim/adamw.py#L210" rel="noopener nofollow ugc">return the loss</a> returned by the closure. Maybe this is useful, but it does not seem commonly used.</p>
</li>
</ol>