How to Penalize Norm of End-to-End Jacobian

RylanSchaeffer: <p>I want to train a network using a modified loss function that has both a typical classification loss (e.g. <code>nn.CrossEntropyLoss</code>) as well as a penalty on the Frobenius norm of the end-to-end Jacobian (i.e. if f(x) is the output of the network, \nabla_x f(x)).</p>
<p>I’ve implemented a model that can successfully learn using <code>nn.CrossEntropyLoss</code>. However, when I try adding the second loss function (by doing two backwards passes), my training loop runs, but the model never learns. Furthermore, if I calculate the end-to-end Jacobian, but don’t include it in the loss function, the model also never learns. At a high level, my code does the following:</p>
<ol>
<li>Forward pass to get predicted classes, <code>yhat</code>, from inputs <code>x</code>
</li>
<li>Call <code>yhat.backward(torch.ones(appropriate shape), retain_graph=True)</code>
</li>
<li>Jacobian norm = <code>x.grad.data.norm(2)</code>
</li>
<li>Set loss equal to classification loss + scalar coefficient * jacobian norm</li>
<li>Run <code>loss.backward()</code>
</li>
</ol>
<p>I suspect that I’m misunderstanding how <code>backward()</code> works when run twice, but I haven’t been able to find any good resources to clarify this.</p>
<p>Too much is required to produce a working example, so I’ve tried to extract the relevant code:</p>
<pre><code class="lang-auto">def train_model(model, train_dataloader, optimizer, loss_fn, device=None):

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model.train()
    train_loss = 0
    correct = 0
    for batch_idx, (batch_input, batch_target) in enumerate(train_dataloader):
        batch_input, batch_target = batch_input.to(device), batch_target.to(device)
        optimizer.zero_grad()
        batch_input.requires_grad_(True)
        model_batch_output = model(batch_input)
        loss = loss_fn(model_output=model_batch_output, model_input=batch_input, model=model, target=batch_target)
        train_loss += loss.item()  # sum up batch loss
        loss.backward()
        optimizer.step()
</code></pre>
<pre><code class="lang-auto">    def end_to_end_jacobian_loss(model_output, model_input):
        model_output.backward(
            torch.ones(*model_output.shape),
            retain_graph=True)
        jacobian = model_input.grad.data
        jacobian_norm = jacobian.norm(2)
        return jacobian_norm
</code></pre>

RylanSchaeffer: <p>I swapped my previous implementation with <code>.backward()</code> to <code>autograd.grad</code> and it apparently works! What’s the difference?</p>
<pre><code class="lang-auto">    def end_to_end_jacobian_loss(model_output, model_input):
        jacobian = autograd.grad(
            outputs=model_output['penultimate_layer'],
            inputs=model_input,
            grad_outputs=torch.ones(*model_output['penultimate_layer'].shape),
            retain_graph=True,
            only_inputs=True)[0]
        jacobian_norm = jacobian.norm(2)
        return jacobian_norm
</code></pre>

albanD: <p>Hi,</p>
<p>You should never need <code>.data</code> anymore. Also keep in mind that <code>.data</code> break the computational graph so no gradients will flow back.</p>
<p>Here you want to compute the gradient in such a way that you can backward through the gradient computation. This is done with the <code>create_graph=True</code> flag to <code>autograd.grad</code>.</p>
<p>Also, since you give a Tensor of ones to grad_outputs, you get the sums of the columns of your Jacobian. Because the autograd only computes a vector Jacobian product.</p>

RylanSchaeffer: <p>How do I obtain the complete matrix Jacobian?</p>

richard: <p>One can compute the full jacobian matrix by doing backward multiple times through the graph. For example, given a  <code>f</code>  that maps R^n to R^m, then we can compute the jacobian via the following:</p>
<pre><code class="lang-auto">def unit_vectors(length):
    result = []
    for i in range(0, length):
        x = torch.zeros(length)
        x[i] = 1
        result.append(x)
    return result

x = torch.randn(3, requires_grad=True)
y = f(x)
result = [torch.autograd.grad(outputs=[y], inputs=[x], grad_outputs=[unit], retain_graph=True)[0] for unit in unit_vectors(y.size(0))]
jacobian = torch.stack(result, dim=0)
</code></pre>
<p>Following up from the questions in <a href="https://github.com/pytorch/pytorch/issues/10223" class="inline-onebox">[feature request] Forward-mode automatic differentiation · Issue #10223 · pytorch/pytorch · GitHub</a>,</p>
<blockquote>
<p><a href="https://github.com/zou3519">@zou3519</a> also, where do you use the  <code>unit_vectors()</code>  function?</p>
</blockquote>
<p><code>unit_vectors</code> is used in the second-to-last line.</p>
<blockquote>
<p>Why are multiple grad calls necessary? Shouldn’t one be sufficient?</p>
</blockquote>
<p>Given a scalar output <code>f_1</code> and an input vector <code>x</code>, autograd gives the vector <code>(\frac{df}{dx_1}, \frac{df}{dx_2}, ..., \frac{df}{dx_n})</code>.</p>
<p>For the full matrix Jacobian we have an output <code>(f_1, f_2, ..., f_m)</code> and we want all <code>\frac{df_i}{dx_j}</code>. One way to get this is to run autograd once on each <code>f_i</code>, which is what is being done in the example above.<br>
.</p>
<p>There might be a way to compress all of this computation into one backward pass; were you thinking of something specific?</p>

RylanSchaeffer: <blockquote>
<p>were you thinking of something specific?</p>
</blockquote>
<p>Yes, I was hoping to penalize the norm of the entire Jacobian, computed in a single backwards pass. I can’t really rely on a solution that scales with the output dimension of the network. Is there no other way?</p>
<p>Why doesn’t autograd compute the full Jacobian ordinarily? Maybe I need to learn more about forward and backward auto-differentiation.</p>

richard: <p>Reverse mode autograd (what we have in pytorch) is capable of computing vector-Jacobian products. That is, given a function <code>f</code>, an input <code>x</code>, and an arbitrary vector <code>v</code>, autograd can tell you <code>v J</code> where J is the Jacobian of <code>f</code> with <code>x</code>.</p>
<p>With multiple vector Jacobian products (If we query autograd for <code>u1 J</code>, <code>u2 J</code>, <code>u3 J</code>, etc where u is a unit vector) we can reconstruct the full Jacobian, which is what the code I wrote above does.</p>

RylanSchaeffer: <p>Touching back on albanD’s point, if I pass a vector of all ones instead of a single one, I receive the sum of all the columns of the Jacobian matrix?</p>

RylanSchaeffer: <p>I now understand. I thought “vector-Jacobian products” meant “we can compute a Jacobian” and “since the Jacobian is a matrix, we can multiply it with vectors.” But this isn’t the case; the two can’t be computed separately currently.</p>
<p>Damn.</p>

richard: <p>Yes, you’d get a sum of all of the columns as <a class="mention" href="/u/alband">@albanD</a> wrote in his reply.</p>

RylanSchaeffer: <p>Ok thank you. Can I ask what the status on forward-mode automatic differentiation in Pytorch is?</p>

richard: <p>I’m not sure what the status of it is but as with all pytorch features if a lot of users want it the dev team takes that into consideration.</p>
<p>Do you have a specific use case for forward mode AD?</p>

RylanSchaeffer: <blockquote>
<p>Do you have a specific use case for forward mode AD?</p>
</blockquote>
<p>I’m not sure what constitutes a specific use case, but this is for a research project. What would you like to know about it?</p>

richard: <p>Oh I am curious how forward-mode AD would enable your research project. Forward mode AD computes Jacobian vector products (as opposed to vector Jacobian products that are computed with reverse-mode AD).</p>

RylanSchaeffer: <p>Maybe forward-mode isn’t what I’m looking for. What mode does jax use to compute the full Jacobian?</p>

richard: <p>I am not sure how jax computes the full Jacobian, but I have always been curious about that. I might do some digging later.</p>
<p>To summarize our discussion above, it sounds like you just want to compute the full Jacobian and use it as a penalty. I posted some code above that computes a Jacobian by invoking autograd.grad multiple times; why doesn’t that work with your use case?</p>

RylanSchaeffer: <blockquote>
<p>why doesn’t that work with your use case?</p>
</blockquote>
<p>I have about 10 different architectures with output dimensions of order 10^3, so running 10 * 1000 backward passes per gradient step is probably far too slow. Unless maybe you think otherwise?</p>

RylanSchaeffer: <blockquote>
<p>I am not sure how jax computes the full Jacobian, but I have always been curious about that. I might do some digging later.</p>
</blockquote>
<p>If you do find out, please let me know <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

RylanSchaeffer: <p>Also, I appreciate you taking the time to help me.</p>

richard: <p>Yeah, that sounds pretty bad. I found this tracking issue for your particular problem (computing a Jacobian efficiently): <a href="https://github.com/pytorch/pytorch/issues/23475">https://github.com/pytorch/pytorch/issues/23475</a> so you can subscribe to that. Unfortunately I don’t have any suggestions for easy ways to speed your case up <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>