8
$\begingroup$

Basic question

Is the some existing mathematical notation to mean "treat this term as a constant when differentiating"? This would be the equivalent of detach in pytorch or stop_gradient in tensorflow and jax.

When I asked this on twitter a helpful suggestion was that $\bot$ (\bot) had been used in one paper (https://arxiv.org/abs/1802.05098), and sg in another (https://arxiv.org/abs/1711.00937). I'll use $\bot$ in the examples below.

I'd like to know if there is a well known notation for this in some literature.

Trivial example for basic question

If $x$ is a function of $t$ we could express $\frac{\partial}{\partial t} f(x,t)$ as $\frac{d}{dt} f(\bot x, t)$.

Extended question

It would be nice to have notation meaning "treat as a constant when differentiating with respect to a particular variable, but otherwise treat as normal" e.g. $\bot_\phi$

Real example covering extended question

In a variational inference derivation I've defined:

$$\mathcal{K}(\phi_1, \phi_2) = E[ \log \pi(\theta_1, \theta_2) + \log f(y_1 | \theta_1) + \log f(y_2 | \theta_1, \theta_2) - \log q_1(\theta_1; \phi_1) - \log q_2(\theta_2 | \theta_1; \phi_2)]. $$

I'd like to calculate $\nabla_{\phi_1} \mathcal{K}$ zeroing out certain terms, but $\nabla_{\phi_2} \mathcal{K}$ without zeroing out anything. It would be nice to write this by defining: $$\mathcal{K}(\phi_1, \phi_2) = E[ \log \pi(\theta_1, \theta_2) + \log f(y_1 | \theta_1) + \bot_{\phi_1}[\log f(y_2 | \theta_1, \theta_2)] - \log q_1(\theta_1; \phi_1) - \bot_{\phi_1}[\log q_2(\theta_2 | \theta_1; \phi_2)]]. $$

For completeness, here's the details of the other notation used in this example:

  • $\pi(\theta_1, \theta_2)$ is a prior density for the parameters.
  • $f(y_1, \theta_1)$ and $f(y_2 | \theta_1, \theta_2)$ are densities for two observations.
  • The generative definition of the approximate posterior is: $$ \theta_1 = g(\epsilon_1, \phi_1), \theta_2 = g(\theta_1, \epsilon_2, \phi_2), \epsilon_1 \sim N(0,I), \epsilon_2 \sim N(0,I). $$ so that the expectation in the defintion of $\mathcal{K}$ is over $\epsilon_1$ and $\epsilon_2$
  • The densities for the approximate posterior are $q(\theta_1; \phi_1), q(\theta_2 | \theta_1; \phi_2)$.
Improve this question
$\endgroup$
4
  • 1
    $\begingroup$ Is the last term in your second definition of $\mathcal{K}$ correct? You have a $\bot_{\phi_1}$ there, but no $\phi_1$ in the term itself, so it is not clear to me what the $\bot_{\phi_1}$ is meant to do here. $\endgroup$
    – jochen
    Commented Apr 14, 2021 at 10:57
  • 1
    $\begingroup$ @jochen Yes this is not immediately obvious. The dependence on $\phi_1$ is because $\theta_2$ depends on $\theta_1$ which depends on $\phi_1$. Perhaps a clearer way to show this would have been to move the $\bot$ into the definition of $\theta_2$ by writing it as $\theta_2 = g(\bot \theta_1, \epsilon_2, \phi_2)$. $\endgroup$
    – Dennis Prangle
    Commented Apr 14, 2021 at 11:01
  • 1
    $\begingroup$ How about using two different symbols for "$\phi$ as usual" and "$\phi$ which is ignored for gradients", for example $\phi$ and $\tilde\phi$. Then you can first define a function $\tilde {K}(\phi_1, \tilde\phi_1, \phi_2)$, $\partial/\partial_{\phi_1} \tilde{K}$ gives the derivative you want, and then ${K}(\phi_1, \phi_2) = \tilde{K}(\phi_1, \phi_1, \phi_2)$ is the function itself. Too clumsy? $\endgroup$
    – jochen
    Commented Apr 14, 2021 at 11:02
  • 1
    $\begingroup$ @jochen Yes this is a nice balance between being formally mathematically correct and simple notationally. I wonder if it illustrates exactly where stop gradient would need to go in the code - I'll need to go back to my derivation and think about this. Definitely a good option! $\endgroup$
    – Dennis Prangle
    Commented Apr 14, 2021 at 11:10

1 Answer 1

Reset to default
4
$\begingroup$

I'm not aware of any standard notation for doing so on a "loss-level", in fact there's no stop_gradient function (in a mathematical sense) that would act as identity, but have a zero derivative. However, one can use standard substitution notation at the "gradient level":

$$ \left(\frac{d}{dt} f(x(s), t)\right) \Bigl|_{s=t} $$

So one way to handle your extended question is to define an extended $\mathcal{K}$:

$$ \begin{align*} \hat{\mathcal{K}}(\phi_1, \phi_2, \varphi_1, \varphi_2) = \mathop{\mathbb{E}}_{\substack{\theta_1 = g(\epsilon_1, \varphi_1) \\ \theta_2 = g(\theta_1, \epsilon_2, \varphi_2)}}[ \log \pi(\theta_1, \theta_2) &+ \log f(y_1 | \theta_1) + \log f(y_2 | \theta_1, \theta_2) \\ &- \log q_1(\theta_1; \phi_1) - \log q_2(\theta_2 | \theta_1; \phi_2)]. \end{align*} $$

Now, the gradient you're interested in is $(\nabla_{\phi_1} \hat{\mathcal{K}}(\phi_1, \phi_2, \varphi_1, \phi_2))|_{\varphi_1 = \phi_1}$

Improve this answer
$\endgroup$
3
  • $\begingroup$ Thanks for the answer! I agree that "stop gradient" is not a standard mathematical operation, and that formally you need to proceed as you've outlined: define two copies of the variable whose gradient you want, fix them as having the same value, but only differentiate with respect to one of them. However I think "stop gradient" would be a useful shorthand notation, especially so it's easy to move between maths and autodiff code without introducing errors. My plan at the moment is to use the shorthand in my main paper, and switch to the formal version for full derivations in the appendix. $\endgroup$
    – Dennis Prangle
    Commented Apr 14, 2021 at 10:58
  • 1
    $\begingroup$ Yes, the proposed notation is definitely too cumbersome and I think properly introduced $\bot_\phi$ is a better practical alternative. $\endgroup$
    – Artem Sobolev
    Commented Apr 14, 2021 at 12:46
  • $\begingroup$ (Note - I'm not writing this up in a paper any more as another group got to this research idea first) $\endgroup$
    – Dennis Prangle
    Commented Nov 29, 2021 at 17:16

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.