Backprop through functional

username: randbit
cooked: <p>I’m trying to backprop through a <a href="https://en.wikipedia.org/wiki/Higher-order_function" rel="noopener nofollow ugc">higher-order function</a> (a function that takes a function as argument), specifically a <a href="https://en.wikipedia.org/wiki/Functional_(mathematics)" rel="noopener nofollow ugc">functional</a> (a higher-order function that returns a scalar). Here is a simple example:</p>
<pre><code class="lang-auto">import torch

class Functional(torch.autograd.Function):
    @staticmethod
    def forward(ctx, f):
        value = f(2)**2 - f(1)
        ctx.save_for_backward(value)
        return value
    @staticmethod
    def backward(ctx, grad_output):
        value = ctx.saved_tensors
        value.backward(grad_output)
        return None

class Function(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.tensor(0.)
    def forward(self, x):
        return self.a * x

function = Function()
functional = Functional.apply

value = functional(function)
value.backward() # RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>
<p>Unfortunately, this throws an error. How can I fix this?</p>
<p>Here is another example, a Monte Carlo estimate of the inner product between two functions on the unit hypercube:</p>
<pre><code class="lang-auto">import torch

dim = 2
samples = 100

class InnerProduct(torch.autograd.Function):
    @staticmethod
    def forward(ctx, f, g):
        x = torch.rand(samples, dim)
        f_x = f(x)
        g_x = g(x)
        ctx.save_for_backward(f_x, g_x)
        return (f_x * g_x).mean(-1)
    @staticmethod
    def backward(ctx, grad_output):
        f_x, g_x = ctx.saved_tensors
        for f_x_i, g_x_i in zip(f_x, g_x):
            f_x_i.backward(g_x_i / samples * grad_output)
            g_x_i.backward(f_x_i / samples * grad_output)
        return None, None
inner_product = InnerProduct.apply

class ScalarFunction(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(dim, 10),
            torch.nn.ELU(),
            torch.nn.Linear(10, 10),
            torch.nn.ELU(),
            torch.nn.Linear(10, 1),
            torch.nn.Flatten(-2, -1)
        )
    def forward(self, x):
        return self.net(x)

f = ScalarFunction()
g = ScalarFunction()
a = inner_product(f, g)
print(a)
a.backward() # RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>

username: tom
cooked: <p>There are a few things between your current code and something that works:</p>
<ul>
<li>the autograd engine actually sees when none of the inputs require a gradient and decides that the entire gradient business can thus be skipped, you could add a dummy argument to avoid that,</li>
<li>torch.autograd.Function.forward you are implicitly in a <code>torch.no_grad</code> context. You need to use <code>with torch.enable_grad():</code> to re-enable gradient tracking.</li>
<li>you didn’t have any gradient requirement in your function, without that, the exercise seems futile (probably you’d want to check and raise an exception or not run backward or something).</li>
<li>but you want to detach the Function.forward return value to prevent bad things to happen™,</li>
<li>I will admit that I have reservations about using <code>.backward</code> in <code>autograd.Functions</code>, but this is just a feeling and I don’t have an alternative.</li>
</ul>
<p>So with these ideas we might produce the following and it runs and puts a gradient into a.grad.</p>
<pre><code class="lang-auto">import torch

class Functional(torch.autograd.Function):
    @staticmethod
    def forward(ctx, f, dummy):
        with torch.enable_grad():
            value = f(2)**2 - f(1)
            ctx.save_for_backward(value)
        return value.detach()
    @staticmethod
    def backward(ctx, grad_output):
        value, = ctx.saved_tensors
        value.backward(grad_output)
        return None, None

class Function(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.tensor(0., requires_grad=True)
    def forward(self, x):
        return self.a * x

dummy = torch.empty((), requires_grad=True)
function = Function()
functional = lambda f: Functional.apply(f, dummy)

value = functional(function)
value.backward() 
</code></pre>
<p>The more important question, of course, is <em>why</em> you would want an autograd.Function if you don’t need the gradients of any inputs. The same result could be obtained if you just used a plain Python function instead of <code>torch.autograd.Functions</code> and without all the complications:</p>
<pre><code class="lang-auto">def functional(f):
  return f(2)**2 - f(1)
</code></pre>
<p>would be a more correct, concise and less indirect way to do things.</p>
<p><code>(lambda: print("Best regards"))()</code></p>
<p>Thomas</p>

username: randbit
cooked: <p><a class="mention" href="/u/tom">@tom</a> The reason I’m doing this is that the forward pass takes a parametrized distribution as input and estimates a property of the distribution by sampling from it. This sampling means I can’t backpropagate directly to the parameters of the distribution. However, I know another estimator for the <em>gradient</em> of the property with respect to the distribution, and this is what I use in the backward pass.</p>

username: randbit
cooked: <p><a class="mention" href="/u/tom">@tom</a>  Is there any way to avoid the dummy input? Perhaps some PyTorch setting to force computation of the backward pass?</p>

username: tom
cooked: <p>You didn’t say what doesn’t work with a plain Python function.</p>

username: randbit
cooked: <p><a class="mention" href="/u/tom">@tom</a> How would the backpropagation phase work through the plain function when the forward pass isn’t differentiable because it requires sampling?</p>

username: tom
cooked: <p>Well, so</p>
<ul>
<li>if you have inputs that are differentiable (and I take it you have), you can use these instead of the dummy parameter,</li>
<li>if you don’t have inputs that are differentiable (which is why you need the dummy), you don’t need an autograd.Function.</li>
</ul>
<p>What I think might make the examples easier is move from passing in a module to passing in a (python) function that doesn’t use global variables instead.<br>
Then you can have the parametrized function and the parameters as part of the inputs in functional and it’ll work with the recipe above. I’m doing something like that (but not with sampling and taking derivatives of the density or so) in my <a href="https://nbviewer.jupyter.org/github/t-vi/pytorch-tvmisc/blob/master/misc/Implicit_Functions_Pytorch.ipynb" rel="noopener nofollow ugc">Implicit Function</a> notebook. (Nowadays you can do the backward that I’m doing in the forward there in the backward, but I did that a looong time ago.)</p>
<p>Best regards</p>
<p>Thomas</p>