Forward Mode AD Example on MNIST

Ryan_Devera: <p>Are there any examples of applying forward mode AD in a real scenario? For example, training a 3-layer MLP on the MNIST data? I am surprised that I can not find anything on this topic,</p>

Ryan_Devera: <p>Even in the documentation in the simple Linear example. I am understanding the output as the computed values of the derivative but this is for the final output how do we get the intermediate calculations as well? I want to assign these to .grad and then use an optimizer:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://pytorch.org/tutorials/intermediate/forward_ad_usage.html">
  <header class="source">

      <a href="https://pytorch.org/tutorials/intermediate/forward_ad_usage.html" target="_blank" rel="noopener nofollow ugc">pytorch.org</a>
  </header>

  <article class="onebox-body">
    

<h3><a href="https://pytorch.org/tutorials/intermediate/forward_ad_usage.html" target="_blank" rel="noopener nofollow ugc">Forward-mode Automatic Differentiation (Beta) — PyTorch Tutorials 2.6.0+cu124...</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


KFrank: <p>Hi Ryan!</p>
<aside class="quote no-group" data-username="Ryan_Devera" data-post="1" data-topic="218588" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ryan_devera/48/53463_2.png" class="avatar"> Ryan_Devera:</div>
<blockquote>
<p>Are there any examples of applying forward mode AD in a real scenario? For example, training a 3-layer MLP on the MNIST data? I am surprised that I can not find anything on this topic,</p>
</blockquote>
</aside>
<p>I am not aware of any (but there could be something somewhere).</p>
<p>The reason is that “real” use cases for forward-mode AD are quite specialized.</p>
<p>In standard backward-mode autograd, a single forward / backward pass computes the<br>
gradient of a single scalar (loss) (or more generally, the vector-jacobian product for a<br>
single vector in the output space) with respect to a bunch of parameters.</p>
<p>On the other hand, a single forward-mode forward pass computes the derivatives of a<br>
bunch of output values with respect to a single upstream value (or more generally, the<br>
directional derivative of the output values with respect to a single direction in upstream<br>
parameter space – that is, the jacobian-vector product for a single vector in upstream<br>
parameter space).</p>
<p>Because you normally want the gradient of a single scalar loss with respect to all of your<br>
trainable parameters, you normally use backward-mode autograd, which gives this to you<br>
with one forward / backward pass.  To get the full gradient with forward-mode autograd,<br>
you would have to perform many forward-mode forward passes, one for each individual<br>
scalar value making up the parameters with respect to which the gradient is being computed.</p>
<p>To train a 3-layer multi-layer perceptron, you would typically compute a single scalar loss<br>
(most likely <code>CrossEntropyLoss</code>) and compute its gradient with respect to all of the model’s<br>
parameters using a single forward / backward pass.  Let’s say your model consists of<br>
1000 individual scalar parameters (these being the elements of the <code>weight</code> and <code>bias</code><br>
tensors of the model’s <code>Linear</code> layers).  You would have to perform 1000 forward-mode<br>
passes to compute the full gradient – a vastly more expensive proposition.</p>
<p>You certainly could use forward mode to compute this gradient – at great cost – and then<br>
follow the rest of the standard workflow, performing gradient-descent optimization to train<br>
the model.</p>
<p>But there is no practical reason to do things this way, so there is little motivation to write<br>
examples that illustrate such a scheme.</p>
<p>Best.</p>
<p>K. Frank</p>