How to use an `optimizer` to update non-differentiable parameters?

MahdiNazemi: <p>Assume there is a non-differentiable <code>nn.Parameter</code> in an equation whose gradient needs to be estimated using a straight-through estimator (STE) before the parameter can be updated. For example,</p>
<p><code>y = 0.5(|x| - |x - alpha| + alpha)</code><br>
<code>y_q = round(y * (2 ** k - 1) / alpha) * alpha / (2 ** k - 1)</code>.</p>
<p>In this equation, <code>alpha</code> is a trainable parameter and the derivative of <code>y_q</code> w.r.t <code>alpha</code> needs to be estimated using an STE.</p>
<p>How can I define the approximated gradient and use it in an <code>optimizer</code> to update the parameter along with the rest of parameters in the network?</p>

tom: <p>Note that it should likely be <code>**</code>, not <code>^</code>.<br>
I like the trick of <code>y_q_diffable = y + (y_q - y).detach()</code>. (In fact, I once proposed to give a lightning talk just on this line of code and whereit is useful.)</p>
<p>I always like to credit <a class="mention" href="/u/hughperkins">@hughperkins</a> for sharing the trick here on the forums when he seen it in a paper and he knows a ton references for applications, too. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Best regards</p>
<p>Thomas</p>

MahdiNazemi: <p>Yes, you’re right. It should be <code>**</code>.</p>
<p>Do you mind explaining this trick a little bit more in detail?</p>

tom: <p>y_q_diffable is y_q ( = y + y_q - y) for the forward. But during backwards, the gradients propagate as if y_q_diffable were y.</p>
<aside class="quote" data-post="10" data-topic="530">
  <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/hughperkins/48/915_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/stop-gradients-for-st-gumbel-softmax/530/10">Stop gradients (for ST gumbel softmax)</a> 
  </div>
  <blockquote>
    From <a href="https://gist.github.com/ericjang/1001afd374c2c3b7752545ce6d9ed349#file-gumbel-softmax-py-L27" rel="noopener nofollow ugc">https://gist.github.com/ericjang/1001afd374c2c3b7752545ce6d9ed349#file-gumbel-softmax-py-L27</a> 
y = tf.stop_gradient(y_hard - y) + y

Whoa, that’s so clever. I had to stare at that for ages before finally figuring that out. So cool <img width="20" height="20" src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title="slight_smile" alt="slight_smile" class="emoji"> . So, the result of this is: 

y is pure one-hot, in terms of value (since we add the soft y, and then subtract it again
the gradients are those of soft y (since all the other terms in this expression have their gradient stripped)
  </blockquote>
</aside>

<p>Best regards</p>
<p>Thomas</p>

gabrieldernbach: <p><a class="mention" href="/u/tom">@tom</a> I am curious, did you prepare that talk in the end? Would you mind sharing some references for applications?</p>

tom: <p>No, I didn’t, but my favourite application that I use in my autograd course is to emulate quantization aware training with it. The course are not freely available, but the particular example is also included in the <a href="https://github.com/t-vi/acdl2020" rel="noopener nofollow ugc">ACDL “Advanced introduction to PyTorch”</a> talk of which I published the slides. I don’t know of any video recording and there wasn’t enough interest to re-record it back then.</p>