How to clip the values of an optimizer?

Johannes_Vogt: <p>Hey,</p>
<p>to prevent NAN values, a common strategy is to use <a href="https://neptune.ai/blog/understanding-gradient-clipping-and-how-it-can-fix-exploding-gradients-problem" rel="noopener nofollow ugc">gradient clipping</a> to cut down all the gradients.</p>
<p>Before the gradient is applied, it is going through an optimizer using e.g. momentum. How can I clip these parameters as well as the final update applied to the weights (efficiently)?</p>
<p>Is there any reason not to clip the optimizer?</p>

ptrblck: <p>I haven’t heard of approaches to clip the internal states of an optimizer, but you could iterate its state and apply your clipping on the desired internal attributes.</p>

qq-me: <p>You can do this:</p>
<pre data-code-wrap="py"><code class="lang-py">import torch

@torch.no_grad
def step_with_update_clip_(optimizer, value=1.0):

    # extract parameters from optimizer
    params: list[torch.Tensor] = [p for g in optimizer.param_groups for p in g['params'] if p.grad is not None]
    params_before = [p.clone() for p in params] # store parameters before step
    optimizer.step() # update parameters

    # update is difference in parameters before and after step, apply clipping to it
    clipped_update = [torch.clip(p - p_before, -value, value) for p, p_before in zip(params, params_before)]

    # revert parameters and add clipped update instead
    for p, p_before, u in zip(params, params_before, clipped_update):
        p.set_(p_before + u)

</code></pre>
<p>Replace <code>optimizer.step()</code> with <code>step_with_update_clip_(optimizer)</code>.</p>
<p>Example of how to use it:</p>
<pre data-code-wrap="py"><code class="lang-py">model = torch.nn.Linear(2,2)
optimizer = torch.optim.Adam(model.parameters(), 1e-2)
criterion = torch.nn.MSELoss()

inputs = torch.randn(10, 2)
targets = torch.randn(10, 2)

for _ in range(1000):
    preds = model(inputs)
    loss = criterion(preds, targets)
    optimizer.zero_grad()
    loss.backward()
    step_with_update_clip_(optimizer, value=1.0)
    print(loss)
</code></pre>
<p>That does require extra clone and subtract operations but they have very little overhead compared to total number of operations performed in the update rule.</p>
<p>Also I am developing a library <a href="https://github.com/inikishev/torchzero" class="inline-onebox" rel="noopener nofollow ugc">GitHub - inikishev/torchzero: Modular optimization library for PyTorch.</a> which does exactly that, you can chain various gradient transformations, for example a typical AdamW with gradient clipping is</p>
<pre data-code-wrap="py"><code class="lang-py">[ClipValue(2), Adam(), WeightDecay(1e-2), LR(1e-3)]
</code></pre>
<p>but you could apply clipping to Adam’s update by moving ClipValue after Adam:</p>
<pre data-code-wrap="py"><code class="lang-py">[Adam(), ClipValue(2), WeightDecay(1e-2), LR(1e-3)]
</code></pre>
<p>And there is no redundant cloning and subtracting.</p>
<p>As of now the library is not fully finished and basically undocumented, so although it is fully useable, I don’t “advertise” it too much, but I am actually quite close to finishing it</p>

Johannes_Vogt: <p>This would solve also another question, of how to set it. Thank you for that!</p>