Recommended way to modify the Adam optimizer?

TriKri: <p>I want to experiment with modifying the <a href="https://arxiv.org/abs/1412.6980" rel="noopener nofollow ugc">Adam</a> optimizer, by for example implementing the modification mentioned in the articles <a href="https://arxiv.org/abs/1711.05101" rel="noopener nofollow ugc">Decoupled Weight Decay Regularization</a> (I know this is already <a href="https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html" rel="noopener nofollow ugc">included in PyTorch</a>), <a href="https://arxiv.org/abs/2302.06675" rel="noopener nofollow ugc">Symbolic Discovery of Optimization Algorithms</a> and <a href="https://arxiv.org/abs/2411.16085" rel="noopener nofollow ugc">Cautious Optimizers: Improving Training with One Line of Code</a>.</p>
<p>When I look at the source code of <a href="https://pytorch.org/docs/stable/_modules/torch/optim/adam.html#Adam" rel="noopener nofollow ugc"><code>torch.optim.Adam</code></a>, it looks like there are two functions that perform the bulk of the computations: <code>_single_tensor_adam</code> and <code>_multi_tensor_adam</code>. However, these are not member functions of <code>torch.optim.Adam</code>, but instead only seem to be called in the function <code>adam</code>, which is in turn called in the method <code>step</code> in <code>torch.optim.Adam</code>.</p>
<p>Is there a recommended way to modify the Adam optimizer? Is it recommended to derive from <code>torch.optim.Adam</code> and override the <code>step</code> function, or is better to create a completely new optimizer from scratch, by deriving from <code>torch.optim.Optimizer</code>? And what is the purpose of the functions <code>adam</code>, <code>_single_tensor_adam</code> and <code>_multi_tensor_adam</code>? Are those something I should implement too?</p>

ptrblck: <aside class="quote no-group" data-username="TriKri" data-post="1" data-topic="214161">
<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/t/9de0a6/48.png" class="avatar"> TriKri:</div>
<blockquote>
<p>it looks like there are two functions that perform the bulk of the computations: <code>_single_tensor_adam</code> and <code>_multi_tensor_adam</code>. However, these are not member functions of <code>torch.optim.Adam</code>, but instead only seem to be called in the function <code>adam</code>, which is in turn called in the method <code>step</code> in <code>torch.optim.Adam</code>.</p>
</blockquote>
</aside>
<p><code>adam</code> is the functional API which is called in <code>torch.optim.Adam.step</code> as you described. <code>_single_tensor_adam</code> is the default implementation, which you could use as a template for your custom implementation. <code>_multi_tensor_adam</code> is a performance optimization, which you could ignore.</p>

TriKri: <p>Thank you, <a class="mention" href="/u/ptrblck">@ptrblck</a>. When you say that I could use <code>_single_tensor_adam</code> as a template for my custom implementation, do you mean that I could override the <code>step</code> method with a method that looks like <code>_single_tensor_adam</code>?</p>

ptrblck: <p>Yes, you can add the body of <code>_single_tensor_adam</code> to your <code>step</code> method directly or also call it via the <code>functional</code> <code>adam</code> call. For the sake of simplicity you could directly reuse <code>_single_tensor_adam</code>’s logic in <code>step</code>.</p>