Why weight sharing of input and output embeddings is not broken by optim.step()?

alessandro_mondin: <p>I am following this PyTorch <a href="https://github.com/jadore801120/attention-is-all-you-need-pytorch" rel="noopener nofollow ugc">implementation</a> of “Attention Is All You Need” and while initialising the Transformer class (<a href="https://github.com/jadore801120/attention-is-all-you-need-pytorch/blob/master/transformer/Models.py#L181" rel="noopener nofollow ugc">here</a> and <a href="https://github.com/jadore801120/attention-is-all-you-need-pytorch/blob/master/transformer/Models.py#L184" rel="noopener nofollow ugc">here</a>) the author sets the <code>encoder_input_embedding.weight = decoder_input_embedding.weight = decoder_output_embedding.weight.</code> However I can’t understand how is possible that this behaviour is not broken by optimiser.step(): what I mean is that, since this layers are in different position within the architecture, I would expect that the errors back-propagated would modify independently such layers. How does PyTorch handle this in its C++ backend?</p>

ptrblck: <p>Reusing the same layer (or generally any parameter) will capture its usage in the computation graph. The backward pass will then calculate the gradients for this layer using its grad outputs as well as forward activations. Each gradient will be accumulated in the <code>.grad</code> attribute of the corresponding parameter.<br>
The optimizer does not have any visibility into the usage of the parameter it was getting and just updates each parameter using the stored <code>.grad</code> attribute and the <code>step()</code> formula.</p>

alessandro_mondin: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="170976">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>The optimizer does not have any visibility into the usage of the parameter it was getting and just updates each parameter using the stored <code>.grad</code> attribute and the <code>step()</code> formula.</p>
</blockquote>
</aside>
<p>Thanks for your response!<br>
In other ways, you mean that by setting for instance <code>encoder_input_embedding.weight = decoder_input_embedding.weight = decoder_output_embedding.weight</code> the gradients accumulated with back propagation will be for each of the above layers the sum of the gradients accumulated of each individual layer?</p>

srishti-git1110: <p>Hi Alessandro,<br>
With parameter sharing, the same tensor (more specifically: the same python object) will be used throughout the calculations in the computation graph. As <a class="mention" href="/u/ptrblck">@ptrblck</a> has explained, the gradient from every branch in the computation graph will be accumulated in the grad attribute of the parameter.</p>
<p>A simple example using <code>torchviz</code> helps understand this better -</p>
<pre data-code-wrap="python"><code class="lang-python">fc1 = nn.Linear(1, 1)
fc2 = nn.Linear(1, 1)

# parameter initialisation
with torch.no_grad():
  fc1.weight = fc2.weight = nn.Parameter(torch.tensor([[2.0]])) # weight sharing
  fc1.bias = nn.Parameter(torch.tensor([2.0]))
  fc2.bias = nn.Parameter(torch.tensor([4.0]))

inp = torch.tensor([1.0])
out = fc1(inp)
out = fc2(out)
make_dot(out).render("out", format="png")
</code></pre>
<p>gives:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/d/9d39ce3bb836393ee3b1d78c99e56d4833101bba.png" data-download-href="https://discuss.pytorch.org/uploads/default/9d39ce3bb836393ee3b1d78c99e56d4833101bba" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/d/9d39ce3bb836393ee3b1d78c99e56d4833101bba_2_314x499.png" alt="image" data-base62-sha1="mqSOgEGfLral2AAwqvQ72TVeJbs" width="314" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/d/9d39ce3bb836393ee3b1d78c99e56d4833101bba_2_314x499.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/d/9d39ce3bb836393ee3b1d78c99e56d4833101bba_2_471x748.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/9/d/9d39ce3bb836393ee3b1d78c99e56d4833101bba.png 2x" data-dominant-color="F1F2F2"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">524×833 28.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>As can be seen in the graph there’s one single node corresponding to the weight of both layers which ensures the calculated gradients match the theoretical values.</p>
<p>As for your question,</p>
<aside class="quote no-group" data-username="alessandro_mondin" data-post="3" data-topic="170976">
<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/alessandro_mondin/48/46669_2.png" class="avatar"> alessandro_mondin:</div>
<blockquote>
<p>the gradients accumulated with back propagation will be for each of the above layers the sum of the gradients accumulated of each individual layer?</p>
</blockquote>
</aside>
<p>If this equivalence holds theoretically, it should be case in PyTorch as well.</p>